From 8910f504d1e4238f2e90de50820891565f95bf6c Mon Sep 17 00:00:00 2001 From: malavhs Date: Thu, 3 Oct 2024 07:15:11 +0000 Subject: [PATCH 01/14] tests: Implement integration tests covering JumpStart PrivateHub workflows --- src/sagemaker/jumpstart/model.py | 2 + tests/integ/sagemaker/jumpstart/conftest.py | 56 +++++- tests/integ/sagemaker/jumpstart/constants.py | 5 + .../jumpstart/model/test_jumpstart_model.py | 3 + .../jumpstart/private_hub/__init__.py | 0 .../jumpstart/private_hub/model/__init__.py | 0 .../model/test_jumpstart_private_hub_model.py | 175 ++++++++++++++++++ .../jumpstart/private_hub/test_hub.py | 40 ++++ .../jumpstart/private_hub/test_hub_content.py | 38 ++++ tests/integ/sagemaker/jumpstart/utils.py | 11 ++ 10 files changed, 328 insertions(+), 2 deletions(-) create mode 100644 tests/integ/sagemaker/jumpstart/private_hub/__init__.py create mode 100644 tests/integ/sagemaker/jumpstart/private_hub/model/__init__.py create mode 100644 tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py create mode 100644 tests/integ/sagemaker/jumpstart/private_hub/test_hub.py create mode 100644 tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index c9354e020b..0dc316dd98 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/src/sagemaker/jumpstart/model.py @@ -1036,6 +1036,7 @@ 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, @@ -1043,6 +1044,7 @@ def _get_deployment_configs( sagemaker_session=self.sagemaker_session, region=self.region, model_version=self.model_version, + hub_arn=self.hub_arn ) deployment_config_metadata = DeploymentConfigMetadata( diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index c7554f3e51..798e39a5e2 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -16,16 +16,26 @@ 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 @@ -33,8 +43,15 @@ 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...") @@ -42,6 +59,7 @@ def _teardown(): 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): diff --git a/tests/integ/sagemaker/jumpstart/constants.py b/tests/integ/sagemaker/jumpstart/constants.py index b839866b1f..bd36f9a369 100644 --- a/tests/integ/sagemaker/jumpstart/constants.py +++ b/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/"), diff --git a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py index 7733041579..c6301440db 100644 --- a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py +++ b/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( diff --git a/tests/integ/sagemaker/jumpstart/private_hub/__init__.py b/tests/integ/sagemaker/jumpstart/private_hub/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/__init__.py b/tests/integ/sagemaker/jumpstart/private_hub/model/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py new file mode 100644 index 0000000000..9b6c2c186f --- /dev/null +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.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 + diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py new file mode 100644 index 0000000000..9beb83d778 --- /dev/null +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.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 \ No newline at end of file diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py new file mode 100644 index 0000000000..3850e5474b --- /dev/null +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.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 diff --git a/tests/integ/sagemaker/jumpstart/utils.py b/tests/integ/sagemaker/jumpstart/utils.py index 0f2fd01572..c75ef5f938 100644 --- a/tests/integ/sagemaker/jumpstart/utils.py +++ b/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 @@ -114,6 +115,16 @@ def get_tabular_data(data_filename: str) -> Tuple[pd.DataFrame, pd.DataFrame]: 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__( From bfeb2c0b9a359bcb0e8dd14d582b5e30fd3bc01f Mon Sep 17 00:00:00 2001 From: malavhs Date: Thu, 3 Oct 2024 07:16:42 +0000 Subject: [PATCH 02/14] linting --- src/sagemaker/jumpstart/model.py | 2 +- tests/integ/sagemaker/jumpstart/conftest.py | 34 +++++++++---------- .../model/test_jumpstart_private_hub_model.py | 33 ++++++++++-------- .../jumpstart/private_hub/test_hub.py | 20 ++++++----- .../jumpstart/private_hub/test_hub_content.py | 24 ++++++------- tests/integ/sagemaker/jumpstart/utils.py | 8 +++-- 6 files changed, 64 insertions(+), 57 deletions(-) diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index 0dc316dd98..3e1b0f8b1e 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/src/sagemaker/jumpstart/model.py @@ -1044,7 +1044,7 @@ def _get_deployment_configs( sagemaker_session=self.sagemaker_session, region=self.region, model_version=self.model_version, - hub_arn=self.hub_arn + hub_arn=self.hub_arn, ) deployment_config_metadata = DeploymentConfigMetadata( diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index 798e39a5e2..2bde884538 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -48,11 +48,14 @@ def _setup(): 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 = 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...") @@ -133,38 +136,35 @@ def _teardown(): # delete private hubs _delete_hubs(sagemaker_session) - + def _delete_hubs(sagemaker_session): - #list Hubs created by PySDK integration tests + # 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']) + 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']) + sagemaker_session.delete_hub(hub["HubName"]) def _delete_hub_contents(sagemaker_session, test_hub_name): - #list hub_contents for the given hub + # 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 + 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']: + # 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_name=test_hub_name, hub_content_type=HubContentType.MODEL_REFERENCE.value, - hub_content_name=models['HubContentName'] + hub_content_name=models["HubContentName"], ) - - @pytest.fixture(scope="session", autouse=True) def setup(request): diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 9b6c2c186f..9717537fcd 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -50,14 +50,16 @@ "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()) + 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) - ) + hub_instance.create_model_reference(model_arn=get_public_hub_model_arn(hub_instance, model)) + def test_jumpstart_hub_model(setup, add_models): @@ -70,7 +72,7 @@ def test_jumpstart_hub_model(setup, add_models): 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] + hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], ) # uses ml.m5.4xlarge instance @@ -78,6 +80,7 @@ def test_jumpstart_hub_model(setup, add_models): 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" @@ -86,7 +89,7 @@ def test_jumpstart_hub_gated_model(setup, add_models): 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] + hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], ) # uses ml.g6.xlarge instance @@ -104,6 +107,7 @@ def test_jumpstart_hub_gated_model(setup, add_models): assert response is not None + def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): model_id = "meta-textgeneration-llama-2-7b" @@ -115,14 +119,14 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): sagemaker_session = get_sm_session() hub_arn = generate_hub_arn_for_init_kwargs( - hub_name=hub_name, region=region, session=sagemaker_session - ) + 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] + hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], ) # uses ml.g5.2xlarge instance @@ -136,7 +140,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): endpoint_name=model.endpoint_name, sagemaker_session=sagemaker_session, tolerate_vulnerable_model=True, - hub_arn=hub_arn + hub_arn=hub_arn, ) payload = { @@ -149,13 +153,13 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): assert response is not None model = JumpStartModel.attach( - predictor.endpoint_name, - sagemaker_session=sagemaker_session, - hub_name=hub_name) + 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" @@ -166,10 +170,9 @@ def test_instatiating_model(setup, add_models): 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] + 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 - diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py index 9beb83d778..14810371cc 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py @@ -14,27 +14,29 @@ HUB_NAME_PREFIX, ) + @pytest.fixture def hub_instance(): - HUB_NAME=f"{HUB_NAME_PREFIX}-{get_test_suite_id()}" + 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 + # 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"], + description="This is a Test Private Hub.", + display_name="malavhs Test hub", + search_keywords=["jumpstart-sdk-integ-test"], ) - #Create Hub Verifications + # Create Hub Verifications assert create_hub_response is not None - #Describe Hub + # Describe Hub hub_description = hub_instance.describe() assert hub_description is not None - #Delete Hub + # Delete Hub delete_hub_response = hub_instance.delete() - assert delete_hub_response is not None \ No newline at end of file + assert delete_hub_response is not None diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py index 3850e5474b..afe09c1e52 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py @@ -6,9 +6,7 @@ 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.utils import get_public_hub_model_arn from tests.integ.sagemaker.jumpstart.constants import ( ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME, ) @@ -18,21 +16,21 @@ 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()) + hub_instance = Hub( + hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], sagemaker_session=get_sm_session() + ) - #Create Model Reference + # Create Model Reference create_model_response = hub_instance.create_model_reference( - model_arn = get_public_hub_model_arn(hub_instance, model_id) + model_arn=get_public_hub_model_arn(hub_instance, model_id) ) - assert create_model_response is not None + assert create_model_response is not None - #Describe Model - describe_model_response = hub_instance.describe_model( - model_name = model_id - ) + # 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 Reference delete_model_response = hub_instance.delete_model_reference(model_name=model_id) - assert delete_model_response is not None + assert delete_model_response is not None diff --git a/tests/integ/sagemaker/jumpstart/utils.py b/tests/integ/sagemaker/jumpstart/utils.py index c75ef5f938..701d106797 100644 --- a/tests/integ/sagemaker/jumpstart/utils.py +++ b/tests/integ/sagemaker/jumpstart/utils.py @@ -115,16 +115,20 @@ def get_tabular_data(data_filename: str) -> Tuple[pd.DataFrame, pd.DataFrame]: 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"]) + 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'] + return models[0]["hub_content_arn"] + class EndpointInvoker: def __init__( From 705ceb9824aedd52b5fe2a089c11ec5f0afb7d89 Mon Sep 17 00:00:00 2001 From: malavhs Date: Thu, 3 Oct 2024 07:25:53 +0000 Subject: [PATCH 03/14] formating --- tests/integ/sagemaker/jumpstart/conftest.py | 9 +++------ .../sagemaker/jumpstart/model/test_jumpstart_model.py | 3 --- .../model/test_jumpstart_private_hub_model.py | 4 ---- tests/integ/sagemaker/jumpstart/private_hub/test_hub.py | 3 --- .../sagemaker/jumpstart/private_hub/test_hub_content.py | 2 -- 5 files changed, 3 insertions(+), 18 deletions(-) diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index 2bde884538..ab08c0590c 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -30,7 +30,6 @@ HubContentType, ) -from sagemaker.jumpstart.constants import JUMPSTART_LOGGER from tests.integ.sagemaker.jumpstart.utils import ( get_test_artifact_bucket, @@ -46,14 +45,15 @@ def _setup(): 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) - describe_hub_response = hub.describe() - JUMPSTART_LOGGER.info(f"Describe Hub {describe_hub_response}") def _teardown(): @@ -62,7 +62,6 @@ def _teardown(): 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) @@ -146,7 +145,6 @@ def _delete_hubs(sagemaker_session): 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"]) @@ -155,7 +153,6 @@ def _delete_hub_contents(sagemaker_session, test_hub_name): 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"]: diff --git a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py index c6301440db..7733041579 100644 --- a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py +++ b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py @@ -56,9 +56,6 @@ 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( diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 9717537fcd..0dac39870a 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -12,12 +12,8 @@ # 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 diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py index 14810371cc..342478b7e7 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py @@ -1,8 +1,5 @@ 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, diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py index afe09c1e52..02d15185a6 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py @@ -1,4 +1,3 @@ -import pytest import os from sagemaker.jumpstart.hub.hub import Hub @@ -10,7 +9,6 @@ from tests.integ.sagemaker.jumpstart.constants import ( ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME, ) -import tests def test_hub_model_reference(setup): From fa7e47c114ddb9b66369549a484e3180af86f000 Mon Sep 17 00:00:00 2001 From: malavhs Date: Thu, 3 Oct 2024 17:58:30 +0000 Subject: [PATCH 04/14] Only delete the pytest session specific test --- tests/integ/sagemaker/jumpstart/conftest.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index ab08c0590c..6273572448 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -139,7 +139,9 @@ def _teardown(): def _delete_hubs(sagemaker_session): # list Hubs created by PySDK integration tests - list_hub_response = sagemaker_session.list_hubs(name_contains=HUB_NAME_PREFIX) + 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: From 7c50ee89f3fcf721b6ae20bd15fa31a0e6b41b66 Mon Sep 17 00:00:00 2001 From: malavhs Date: Thu, 3 Oct 2024 22:11:54 +0000 Subject: [PATCH 05/14] change scope to session --- .../private_hub/model/test_jumpstart_private_hub_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 0dac39870a..be78672de2 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -47,7 +47,7 @@ } -@pytest.fixture(scope="module") +@pytest.fixture(scope="session") def add_models(): # Create Model References to test in Hub hub_instance = Hub( From cb5f1c709faf5b3095da039f744883bc299e920b Mon Sep 17 00:00:00 2001 From: malavhs Date: Thu, 3 Oct 2024 22:14:32 +0000 Subject: [PATCH 06/14] address nits --- tests/integ/sagemaker/jumpstart/private_hub/test_hub.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py index 342478b7e7..179b809024 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py @@ -23,7 +23,7 @@ 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", + display_name="PySDK integration tests Hub", search_keywords=["jumpstart-sdk-integ-test"], ) From 4bd94ecbca5991ded66bf65b1b73e304b877e65d Mon Sep 17 00:00:00 2001 From: malavhs Date: Fri, 4 Oct 2024 02:24:45 +0000 Subject: [PATCH 07/14] Address test failures --- .../sagemaker/jumpstart/private_hub/test_hub.py | 13 +++++++++++++ .../jumpstart/private_hub/test_hub_content.py | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py index 179b809024..2bccb96524 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub.py @@ -1,3 +1,16 @@ +# 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 pytest from sagemaker.jumpstart.hub.hub import Hub diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py index 02d15185a6..06a2540ae1 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py @@ -1,3 +1,16 @@ +# 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 from sagemaker.jumpstart.hub.hub import Hub From 3fed9f47821784b990362abb0c4ca3429240ff9c Mon Sep 17 00:00:00 2001 From: malavhs Date: Fri, 4 Oct 2024 03:57:24 +0000 Subject: [PATCH 08/14] address typo --- tests/integ/sagemaker/jumpstart/conftest.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index 6273572448..c1ba01ba1e 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -42,11 +42,11 @@ def _setup(): print("Setting up...") - test_suit_id = get_test_suite_id() - test_hub_name = f"{HUB_NAME_PREFIX}{test_suit_id}" + test_suite_id = get_test_suite_id() + test_hub_name = f"{HUB_NAME_PREFIX}{test_suite_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_SUITE_ID: test_suite_id}) os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME: test_hub_name}) # Create a private hub to use for the test session From 6456883015e7ea4bf30662df44e8f5c9057f82df Mon Sep 17 00:00:00 2001 From: malavhs Date: Mon, 7 Oct 2024 00:20:37 +0000 Subject: [PATCH 09/14] address comments --- src/sagemaker/jumpstart/constants.py | 2 - src/sagemaker/jumpstart/model.py | 1 + tests/integ/sagemaker/jumpstart/conftest.py | 5 +-- tests/integ/sagemaker/jumpstart/constants.py | 2 - .../jumpstart/model/test_jumpstart_model.py | 2 +- .../model/test_jumpstart_private_hub_model.py | 38 ++++++++----------- .../jumpstart/private_hub/test_hub_content.py | 3 -- tests/integ/sagemaker/jumpstart/utils.py | 5 --- 8 files changed, 19 insertions(+), 39 deletions(-) diff --git a/src/sagemaker/jumpstart/constants.py b/src/sagemaker/jumpstart/constants.py index 03a0ebe545..eff7f53717 100644 --- a/src/sagemaker/jumpstart/constants.py +++ b/src/sagemaker/jumpstart/constants.py @@ -222,8 +222,6 @@ JUMPSTART_MODEL_HUB_NAME = "SageMakerPublicHub" -JUMPSTART_MODEL_HUB_NAME = "SageMakerPublicHub" - JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json" JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY = "proprietary-sdk-manifest.json" diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py index 3e1b0f8b1e..e6952b2154 100644 --- a/src/sagemaker/jumpstart/model.py +++ b/src/sagemaker/jumpstart/model.py @@ -447,6 +447,7 @@ def retrieve_example_payload(self) -> JumpStartSerializablePayload: return payloads.retrieve_example( model_id=self.model_id, model_version=self.model_version, + hub_arn=self.hub_arn, model_type=self.model_type, region=self.region, tolerate_deprecated_model=self.tolerate_deprecated_model, diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index c1ba01ba1e..14edd8fd54 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -23,7 +23,6 @@ ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME, HUB_NAME_PREFIX, JUMPSTART_TAG, - SM_JUMPSTART_PUBLIC_HUB_NAME, ) from sagemaker.jumpstart.types import ( @@ -37,7 +36,7 @@ get_sm_session, ) -from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME +from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME, JUMPSTART_MODEL_HUB_NAME def _setup(): @@ -144,7 +143,7 @@ def _delete_hubs(sagemaker_session): ) for hub in list_hub_response["HubSummaries"]: - if hub["HubName"] != SM_JUMPSTART_PUBLIC_HUB_NAME: + if hub["HubName"] != JUMPSTART_MODEL_HUB_NAME: # delete all hub contents first _delete_hub_contents(sagemaker_session, hub["HubName"]) sagemaker_session.delete_hub(hub["HubName"]) diff --git a/tests/integ/sagemaker/jumpstart/constants.py b/tests/integ/sagemaker/jumpstart/constants.py index bd36f9a369..27279feecf 100644 --- a/tests/integ/sagemaker/jumpstart/constants.py +++ b/tests/integ/sagemaker/jumpstart/constants.py @@ -41,8 +41,6 @@ def _to_s3_path(filename: str, s3_prefix: Optional[str]) -> str: JUMPSTART_TAG = "JumpStart-SDK-Integ-Test-Suite-Id" -SM_JUMPSTART_PUBLIC_HUB_NAME = "SageMakerPublicHub" - HUB_NAME_PREFIX = "PySDK-HubTest-" TRAINING_DATASET_MODEL_DICT = { diff --git a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py index 7733041579..ec98786da4 100644 --- a/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py +++ b/tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py @@ -229,7 +229,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup): @mock.patch("sagemaker.jumpstart.cache.JUMPSTART_LOGGER.warning") -def test_instatiating_model(mock_warning_logger, setup): +def test_instantiating_model(mock_warning_logger, setup): model_id = "catboost-regression-model" diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index be78672de2..9fe9da34a4 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -48,7 +48,7 @@ @pytest.fixture(scope="session") -def add_models(): +def add_model_references(): # 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() @@ -57,27 +57,27 @@ def add_models(): 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()}") +def test_jumpstart_hub_model(setup, add_model_references): model_id = "catboost-classification-model" + sagemaker_session = get_sm_session() + model = JumpStartModel( model_id=model_id, - role=get_sm_session().get_caller_identity_arn(), - sagemaker_session=get_sm_session(), + role=sagemaker_session.get_caller_identity_arn(), + sagemaker_session=sagemaker_session, hub_name=os.environ[ENV_VAR_JUMPSTART_SDK_TEST_HUB_NAME], ) - # uses ml.m5.4xlarge instance - model.deploy( + predictor = model.deploy( tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}], ) + assert sagemaker_session.endpoint_in_service_or_not(predictor.endpoint_name) -def test_jumpstart_hub_gated_model(setup, add_models): + +def test_jumpstart_hub_gated_model(setup, add_model_references): model_id = "meta-textgeneration-llama-3-2-1b" @@ -88,23 +88,19 @@ def test_jumpstart_hub_gated_model(setup, add_models): 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}, - } + payload = model.retrieve_example_payload() - response = predictor.predict(payload, custom_attributes="accept_eula=true") + response = predictor.predict(payload) assert response is not None -def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): +def test_jumpstart_gated_model_inference_component_enabled(setup, add_model_references): model_id = "meta-textgeneration-llama-2-7b" @@ -125,7 +121,6 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): 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, @@ -139,10 +134,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): hub_arn=hub_arn, ) - payload = { - "inputs": "some-payload", - "parameters": {"max_new_tokens": 256, "top_p": 0.9, "temperature": 0.6}, - } + payload = model.retrieve_example_payload() response = predictor.predict(payload) @@ -156,7 +148,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup, add_models): assert model.inference_component_name == predictor.component_name -def test_instatiating_model(setup, add_models): +def test_instantiating_model(setup, add_model_references): model_id = "catboost-regression-model" diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py index 06a2540ae1..ca1164b52a 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py @@ -31,17 +31,14 @@ def test_hub_model_reference(setup): 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 diff --git a/tests/integ/sagemaker/jumpstart/utils.py b/tests/integ/sagemaker/jumpstart/utils.py index 701d106797..b16c147e92 100644 --- a/tests/integ/sagemaker/jumpstart/utils.py +++ b/tests/integ/sagemaker/jumpstart/utils.py @@ -121,11 +121,6 @@ def get_public_hub_model_arn(hub: Hub, model_id: str) -> str: 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"] From bac00dd2f1b7b9b4fe1fc9625225b15cefeae59d Mon Sep 17 00:00:00 2001 From: malavhs Date: Mon, 7 Oct 2024 03:25:34 +0000 Subject: [PATCH 10/14] resolve flake8 errors --- .../private_hub/model/test_jumpstart_private_hub_model.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 9fe9da34a4..bb7f436a39 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -21,8 +21,6 @@ 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 556d120a78f7499ae2555c389b53f393c07a7d1a Mon Sep 17 00:00:00 2001 From: malavhs Date: Mon, 7 Oct 2024 05:03:10 +0000 Subject: [PATCH 11/14] implement throttle handling --- tests/integ/sagemaker/jumpstart/conftest.py | 39 ++++++++----------- .../model/test_jumpstart_private_hub_model.py | 11 +++++- tests/integ/sagemaker/jumpstart/utils.py | 29 ++++++++++++++ 3 files changed, 56 insertions(+), 23 deletions(-) diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index 14edd8fd54..24a18694b3 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -34,6 +34,7 @@ get_test_artifact_bucket, get_test_suite_id, get_sm_session, + with_exponential_backoff, ) from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME, JUMPSTART_MODEL_HUB_NAME @@ -62,6 +63,8 @@ def _teardown(): 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) sagemaker_client = boto3_session.client( @@ -133,36 +136,28 @@ def _teardown(): bucket.objects.filter(Prefix=test_suite_id + "/").delete() # delete private hubs - _delete_hubs(sagemaker_session) + _delete_hubs(sagemaker_session, test_hub_name) -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] +def _delete_hubs(sagemaker_session, hub_name): + # list and delete all hub contents first + list_hub_content_response = sagemaker_session.list_hub_contents( + hub_name=hub_name, hub_content_type=HubContentType.MODEL_REFERENCE.value ) + for model in list_hub_content_response["HubContentSummaries"]: + _delete_hub_contents(sagemaker_session, hub_name, model) - for hub in list_hub_response["HubSummaries"]: - if hub["HubName"] != JUMPSTART_MODEL_HUB_NAME: - # delete all hub contents first - _delete_hub_contents(sagemaker_session, hub["HubName"]) - sagemaker_session.delete_hub(hub["HubName"]) + sagemaker_session.delete_hub(hub_name) -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 +@with_exponential_backoff() +def _delete_hub_contents(sagemaker_session, hub_name, model): + sagemaker_session.delete_hub_content_reference( + hub_name=hub_name, + hub_content_type=HubContentType.MODEL_REFERENCE.value, + hub_content_name=model["HubContentName"], ) - # 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): diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index bb7f436a39..03268056b5 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -13,6 +13,7 @@ from __future__ import absolute_import import os +import random import time import pytest @@ -20,6 +21,7 @@ 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 botocore.exceptions import ClientError import tests.integ @@ -32,6 +34,7 @@ from tests.integ.sagemaker.jumpstart.utils import ( get_public_hub_model_arn, get_sm_session, + with_exponential_backoff, ) MAX_INIT_TIME_SECONDS = 5 @@ -45,6 +48,11 @@ } +@with_exponential_backoff() +def create_model_reference(hub_instance, model_arn): + hub_instance.create_model_reference(model_arn=model_arn) + + @pytest.fixture(scope="session") def add_model_references(): # Create Model References to test in Hub @@ -52,7 +60,8 @@ def add_model_references(): 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)) + model_arn = get_public_hub_model_arn(hub_instance, model) + create_model_reference(hub_instance, model_arn) def test_jumpstart_hub_model(setup, add_model_references): diff --git a/tests/integ/sagemaker/jumpstart/utils.py b/tests/integ/sagemaker/jumpstart/utils.py index b16c147e92..47dc1f45d3 100644 --- a/tests/integ/sagemaker/jumpstart/utils.py +++ b/tests/integ/sagemaker/jumpstart/utils.py @@ -14,6 +14,8 @@ import functools import json +import random +import time import uuid from typing import Any, Dict, List, Tuple import boto3 @@ -21,6 +23,7 @@ import os from botocore.config import Config +from botocore.exceptions import ClientError import pytest @@ -125,6 +128,32 @@ def get_public_hub_model_arn(hub: Hub, model_id: str) -> str: return models[0]["hub_content_arn"] +def with_exponential_backoff(max_retries=5, initial_delay=1, max_delay=60): + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + retries = 0 + while True: + try: + return func(*args, **kwargs) + except ClientError as e: + if retries >= max_retries or e.response["Error"]["Code"] not in [ + "ThrottlingException", + "TooManyRequestsException", + ]: + raise + delay = min(initial_delay * (2**retries) + random.random(), max_delay) + print( + f"Retrying {func.__name__} in {delay:.2f} seconds... (Attempt {retries + 1}/{max_retries})" + ) + time.sleep(delay) + retries += 1 + + return wrapper + + return decorator + + class EndpointInvoker: def __init__( self, From 8ff04d3b41068c2973173b9ca0ccc0d4ab30b9fc Mon Sep 17 00:00:00 2001 From: malavhs Date: Mon, 7 Oct 2024 05:06:21 +0000 Subject: [PATCH 12/14] flake8 --- tests/integ/sagemaker/jumpstart/conftest.py | 2 -- .../private_hub/model/test_jumpstart_private_hub_model.py | 1 - 2 files changed, 3 deletions(-) diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index 24a18694b3..d8cadcc513 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -37,8 +37,6 @@ with_exponential_backoff, ) -from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME, JUMPSTART_MODEL_HUB_NAME - def _setup(): print("Setting up...") diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 03268056b5..98b1ba6d10 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -21,7 +21,6 @@ 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 botocore.exceptions import ClientError import tests.integ From d79b8a310ec705057410219b010bb7eeb135c493 Mon Sep 17 00:00:00 2001 From: malavhs Date: Mon, 7 Oct 2024 05:23:34 +0000 Subject: [PATCH 13/14] flake8 --- tests/integ/sagemaker/jumpstart/conftest.py | 1 + .../private_hub/model/test_jumpstart_private_hub_model.py | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integ/sagemaker/jumpstart/conftest.py b/tests/integ/sagemaker/jumpstart/conftest.py index d8cadcc513..260b0f2b22 100644 --- a/tests/integ/sagemaker/jumpstart/conftest.py +++ b/tests/integ/sagemaker/jumpstart/conftest.py @@ -16,6 +16,7 @@ import boto3 import pytest from botocore.config import Config +from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME from sagemaker.jumpstart.hub.hub import Hub from sagemaker.session import Session from tests.integ.sagemaker.jumpstart.constants import ( diff --git a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py index 98b1ba6d10..751162d2e6 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/model/test_jumpstart_private_hub_model.py @@ -13,7 +13,6 @@ from __future__ import absolute_import import os -import random import time import pytest From e0b84677ceb61a5640d101c312f9ca3409e07c89 Mon Sep 17 00:00:00 2001 From: malavhs Date: Tue, 8 Oct 2024 00:23:49 +0000 Subject: [PATCH 14/14] Adding more assertions --- tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py index ca1164b52a..b25cff2d62 100644 --- a/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py +++ b/tests/integ/sagemaker/jumpstart/private_hub/test_hub_content.py @@ -39,6 +39,8 @@ def test_hub_model_reference(setup): describe_model_response = hub_instance.describe_model(model_name=model_id) assert describe_model_response is not None assert type(describe_model_response) == DescribeHubContentResponse + assert describe_model_response.hub_content_name == model_id + assert describe_model_response.hub_content_type == "ModelReference" delete_model_response = hub_instance.delete_model_reference(model_name=model_id) assert delete_model_response is not None