feat: Integ tests for jumpstart model and estimator (#2865)

Author: evakravi

src/sagemaker/model.py

@@ -406,7 +406,10 @@ def prepare_container_def(
             self._upload_code(deploy_key_prefix, repack=is_repack)
             deploy_env.update(self._script_mode_env_vars())
         return sagemaker.container_def(
-            self.image_uri, self.model_data, deploy_env, image_config=self.image_config
+            self.image_uri,
+            self.repacked_model_data or self.model_data,
+            deploy_env,
+            image_config=self.image_config,
         )
 
     def _upload_code(self, key_prefix: str, repack: bool = False) -> None:

tests/integ/sagemaker/jumpstart/retrieve_uri/conftest.py renamed to tests/integ/sagemaker/jumpstart/conftest.py

@@ -16,16 +16,16 @@
 import boto3
 import pytest
 from botocore.config import Config
+from tests.integ.sagemaker.jumpstart.constants import (
+    ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
+    JUMPSTART_TAG,
+)
 
 
-from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
+from tests.integ.sagemaker.jumpstart.utils import (
     get_test_artifact_bucket,
     get_test_suite_id,
 )
-from tests.integ.sagemaker.jumpstart.retrieve_uri.constants import (
-    ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
-    JUMPSTART_TAG,
-)
 
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
 

tests/integ/sagemaker/jumpstart/retrieve_uri/constants.py renamed to tests/integ/sagemaker/jumpstart/constants.py

@@ -12,25 +12,21 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import json
 import time
-from typing import Any, Dict, List
 import boto3
 import os
 from botocore.config import Config
-import pandas as pd
 
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
-from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
+from tests.integ.sagemaker.jumpstart.utils import (
     get_test_artifact_bucket,
     get_sm_session,
 )
 
 from sagemaker.utils import repack_model
-from tests.integ.sagemaker.jumpstart.retrieve_uri.constants import (
+from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
     JUMPSTART_TAG,
-    ContentType,
 )
 
 
@@ -189,43 +185,3 @@ def create_endpoint(self) -> None:
                 }
             ],
         )
-
-
-class EndpointInvoker:
-    def __init__(
-        self,
-        endpoint_name,
-        region=JUMPSTART_DEFAULT_REGION_NAME,
-        boto_config=Config(retries={"max_attempts": 10, "mode": "standard"}),
-    ) -> None:
-        self.endpoint_name = endpoint_name
-        self.region = region
-        self.config = boto_config
-        self.sagemaker_runtime_client = self.get_sagemaker_runtime_client()
-
-    def _invoke_endpoint(
-        self,
-        body: Any,
-        content_type: ContentType,
-    ) -> Dict[str, Any]:
-        response = self.sagemaker_runtime_client.invoke_endpoint(
-            EndpointName=self.endpoint_name, ContentType=content_type.value, Body=body
-        )
-        return json.loads(response["Body"].read())
-
-    def invoke_tabular_endpoint(self, data: pd.DataFrame) -> Dict[str, Any]:
-        return self._invoke_endpoint(
-            body=data.to_csv(header=False, index=False).encode("utf-8"),
-            content_type=ContentType.TEXT_CSV,
-        )
-
-    def invoke_spc_endpoint(self, text: List[str]) -> Dict[str, Any]:
-        return self._invoke_endpoint(
-            body=json.dumps(text).encode("utf-8"),
-            content_type=ContentType.LIST_TEXT,
-        )
-
-    def get_sagemaker_runtime_client(self) -> boto3.client:
-        return boto3.client(
-            service_name="runtime.sagemaker", config=self.config, region_name=self.region
-        )

tests/integ/sagemaker/jumpstart/retrieve_uri/inference.py

@@ -14,18 +14,18 @@
 
 
 from tests.integ.sagemaker.jumpstart.retrieve_uri.inference import (
-    EndpointInvoker,
     InferenceJobLauncher,
 )
 from sagemaker import environment_variables, image_uris
 from sagemaker import script_uris
 from sagemaker import model_uris
 
-from tests.integ.sagemaker.jumpstart.retrieve_uri.constants import InferenceTabularDataname
+from tests.integ.sagemaker.jumpstart.constants import InferenceTabularDataname
 
-from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
+from tests.integ.sagemaker.jumpstart.utils import (
     download_inference_assets,
     get_tabular_data,
+    EndpointInvoker,
 )
 
 

tests/integ/sagemaker/jumpstart/retrieve_uri/test_inference.py

@@ -12,15 +12,15 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import pandas as pd
-
 
 from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
     get_model_tarball_full_uri_from_base_uri,
+)
+from tests.integ.sagemaker.jumpstart.utils import (
     get_training_dataset_for_model_and_version,
+    EndpointInvoker,
 )
 from tests.integ.sagemaker.jumpstart.retrieve_uri.inference import (
-    EndpointInvoker,
     InferenceJobLauncher,
 )
 from tests.integ.sagemaker.jumpstart.retrieve_uri.training import TrainingJobLauncher
@@ -59,6 +59,8 @@ def test_jumpstart_transfer_learning_retrieve_functions(setup):
         model_id=model_id, model_version=model_version, include_container_hyperparameters=True
     )
 
+    default_hyperparameters["epochs"] = "1"
+
     training_job = TrainingJobLauncher(
         image_uri=image_uri,
         script_uri=script_uri,
@@ -110,10 +112,5 @@ def test_jumpstart_transfer_learning_retrieve_functions(setup):
     )
 
     response = endpoint_invoker.invoke_spc_endpoint(["hello", "world"])
-    entail, no_entail = response[0][0], response[0][1]
-
-    assert entail is not None
-    assert no_entail is not None
 
-    assert pd.isna(entail) is False
-    assert pd.isna(no_entail) is False
+    assert response is not None

tests/integ/sagemaker/jumpstart/retrieve_uri/test_transfer_learning.py

@@ -18,14 +18,16 @@
 from botocore.config import Config
 
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
-from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
-    get_full_hyperparameters,
+from tests.integ.sagemaker.jumpstart.utils import (
     get_test_artifact_bucket,
     get_sm_session,
 )
+from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
+    get_full_hyperparameters,
+)
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 
-from tests.integ.sagemaker.jumpstart.retrieve_uri.constants import (
+from tests.integ.sagemaker.jumpstart.constants import (
     ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
 )
 

tests/integ/sagemaker/jumpstart/retrieve_uri/training.py

@@ -12,26 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import uuid
-from typing import Tuple
-import boto3
-import pandas as pd
-import os
-
-from tests.integ.sagemaker.jumpstart.retrieve_uri.constants import (
-    TEST_ASSETS_SPECS,
-    TMP_DIRECTORY_PATH,
-    TRAINING_DATASET_MODEL_DICT,
-)
-from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
-from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
-
 from sagemaker.s3 import parse_s3_url
-from sagemaker.session import Session
-
-
-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_model_tarball_full_uri_from_base_uri(base_uri: str, training_job_name: str) -> str:
@@ -56,46 +37,3 @@ def get_full_hyperparameters(
         "model-artifact-bucket": bucket,
         "model-artifact-key": key,
     }
-
-
-def get_training_dataset_for_model_and_version(model_id: str, version: str) -> dict:
-    return TRAINING_DATASET_MODEL_DICT[(model_id, version)]
-
-
-def get_sm_session() -> Session:
-    return Session(boto_session=boto3.Session(region_name=JUMPSTART_DEFAULT_REGION_NAME))
-
-
-def get_test_artifact_bucket() -> str:
-    bucket_name = get_sm_session().default_bucket()
-    return bucket_name
-
-
-def download_inference_assets():
-
-    if not os.path.exists(TMP_DIRECTORY_PATH):
-        os.makedirs(TMP_DIRECTORY_PATH)
-
-    for asset, s3_key in TEST_ASSETS_SPECS.items():
-        file_path = os.path.join(TMP_DIRECTORY_PATH, str(asset.value))
-        if not os.path.exists(file_path):
-            download_file(
-                file_path,
-                get_jumpstart_content_bucket(JUMPSTART_DEFAULT_REGION_NAME),
-                s3_key,
-                boto3.client("s3"),
-            )
-
-
-def get_tabular_data(data_filename: str) -> Tuple[pd.DataFrame, pd.DataFrame]:
-
-    asset_file_path = os.path.join(TMP_DIRECTORY_PATH, data_filename)
-
-    test_data = pd.read_csv(asset_file_path, header=None)
-    label, features = test_data.iloc[:, :1], test_data.iloc[:, 1:]
-
-    return label, features
-
-
-def get_test_suite_id() -> str:
-    return str(uuid.uuid4())

tests/integ/sagemaker/jumpstart/retrieve_uri/utils.py

@@ -0,0 +1,81 @@
+# 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 import image_uris, model_uris, script_uris
+from sagemaker.jumpstart.constants import INFERENCE_ENTRYPOINT_SCRIPT_NAME
+from sagemaker.model import Model
+from tests.integ.sagemaker.jumpstart.constants import (
+    ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
+    JUMPSTART_TAG,
+    InferenceTabularDataname,
+)
+from tests.integ.sagemaker.jumpstart.utils import (
+    EndpointInvoker,
+    download_inference_assets,
+    get_sm_session,
+    get_tabular_data,
+)
+
+
+def test_jumpstart_inference_model_class(setup):
+
+    model_id, model_version = "catboost-classification-model", "1.0.0"
+    instance_type, instance_count = "ml.m5.xlarge", 1
+
+    print("Starting inference...")
+
+    image_uri = image_uris.retrieve(
+        region=None,
+        framework=None,
+        image_scope="inference",
+        model_id=model_id,
+        model_version=model_version,
+        instance_type=instance_type,
+    )
+
+    script_uri = script_uris.retrieve(
+        model_id=model_id, model_version=model_version, script_scope="inference"
+    )
+
+    model_uri = model_uris.retrieve(
+        model_id=model_id, model_version=model_version, model_scope="inference"
+    )
+
+    model = Model(
+        image_uri=image_uri,
+        model_data=model_uri,
+        source_dir=script_uri,
+        entry_point=INFERENCE_ENTRYPOINT_SCRIPT_NAME,
+        role=get_sm_session().get_caller_identity_arn(),
+        sagemaker_session=get_sm_session(),
+        enable_network_isolation=True,
+    )
+
+    model.deploy(
+        initial_instance_count=instance_count,
+        instance_type=instance_type,
+        tags=[{"Key": JUMPSTART_TAG, "Value": os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]}],
+    )
+
+    endpoint_invoker = EndpointInvoker(
+        endpoint_name=model.endpoint_name,
+    )
+
+    download_inference_assets()
+    ground_truth_label, features = get_tabular_data(InferenceTabularDataname.MULTICLASS)
+
+    response = endpoint_invoker.invoke_tabular_endpoint(features)
+
+    assert response is not None