address comments

Author: malavhs

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"
 

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,

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"])

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 = {

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"
 

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"
 

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

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"]