feat: integration tests for jumpstart sdk retrieve functions

Author: evakravi

.gitignore

@@ -27,4 +27,5 @@ venv/
 *.swp
 .docker/
 env/
-.vscode/
+.vscode/
+**/tmp

setup.py

@@ -79,6 +79,7 @@ def read_version():
         "fabric>=2.0",
         "requests>=2.20.0, <3",
         "sagemaker-experiments",
+        "regex",
     ],
 )
 

tests/integ/sagemaker/jumpstart/__init__.py

@@ -0,0 +1,89 @@
+# 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 boto3
+import pytest
+
+from tests.integ.sagemaker.jumpstart.retrieve_uri.utils import (
+    get_test_cache_bucket,
+    get_test_suite_id,
+)
+from tests.integ.sagemaker.jumpstart.retrieve_uri.constants import (
+    ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID,
+    JUMPSTART_TAG,
+)
+import os
+
+
+def _setup():
+    print("Setting up...")
+    os.environ.update({ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID: get_test_suite_id()})
+
+
+def _teardown():
+    print("Tearing down...")
+
+    test_cache_bucket = get_test_cache_bucket()
+
+    test_suite_id = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]
+
+    sagemaker_client = boto3.client("sagemaker")
+
+    search_endpoints_result = sagemaker_client.search(
+        Resource="Endpoint",
+        SearchExpression={
+            "Filters": [
+                {"Name": f"Tags.{JUMPSTART_TAG}", "Operator": "Equals", "Value": test_suite_id}
+            ]
+        },
+    )
+
+    endpoint_names = [
+        endpoint_info["Endpoint"]["EndpointName"]
+        for endpoint_info in search_endpoints_result["Results"]
+    ]
+    endpoint_config_names = [
+        endpoint_info["Endpoint"]["EndpointConfigName"]
+        for endpoint_info in search_endpoints_result["Results"]
+    ]
+    model_names = [
+        sagemaker_client.describe_endpoint_config(EndpointConfigName=endpoint_config_name)[
+            "ProductionVariants"
+        ][0]["ModelName"]
+        for endpoint_config_name in endpoint_config_names
+    ]
+
+    # delete test-suite-tagged endpoints
+    for endpoint_name in endpoint_names:
+        sagemaker_client.delete_endpoint(EndpointName=endpoint_name)
+
+    # delete endpoint configs for test-suite-tagged endpoints
+    for endpoint_config_name in endpoint_config_names:
+        sagemaker_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)
+
+    # delete models for test-suite-tagged endpoints
+    for model_name in model_names:
+        sagemaker_client.delete_model(ModelName=model_name)
+
+    # delete test artifact/cache s3 folder
+    s3_resource = boto3.resource("s3")
+    bucket = s3_resource.Bucket(test_cache_bucket)
+    bucket.objects.filter(Prefix=test_suite_id + "/").delete()
+
+
+@pytest.fixture(scope="session", autouse=True)
+def setup(request):
+    _setup()
+
+    request.addfinalizer(_teardown)

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

@@ -0,0 +1,118 @@
+# 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
+
+from enum import Enum
+from typing import Dict
+from typing import Optional
+from typing import Union
+import os
+
+
+def _to_s3_path(filename: str, s3_folder: Optional[str]) -> str:
+    return filename if not s3_folder else f"{s3_folder}/{filename}"
+
+
+_NB_ASSETS_S3_FOLDER = "inference-notebook-assets"
+_TF_FLOWERS_S3_FOLDER = "training-datasets/tf_flowers"
+
+TMP_DIRECTORY_PATH = os.path.join(
+    os.path.abspath(os.path.join(os.path.abspath(__file__), os.pardir)), "tmp"
+)
+
+ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID = "JUMPSTART_SDK_TEST_SUITE_ID"
+
+JUMPSTART_TAG = "JumpStart-SDK-Integ-Test-Suite-Id"
+
+HYPERPARAMETER_MODEL_DICT = {
+    ("huggingface-spc-bert-base-cased", "*"): {
+        "epochs": "1",
+        "adam-learning-rate": "2e-05",
+        "batch-size": "8",
+        "sagemaker_submit_directory": "/opt/ml/input/data/code/sourcedir.tar.gz",
+        "sagemaker_program": "transfer_learning.py",
+        "sagemaker_container_log_level": "20",
+    },
+}
+
+TRAINING_DATASET_MODEL_DICT = {
+    ("huggingface-spc-bert-base-cased", "*"): ("training-datasets/QNLI/"),
+}
+
+
+class ContentType(str, Enum):
+    """Possible value for content type argument of SageMakerRuntime.invokeEndpoint."""
+
+    X_IMAGE = "application/x-image"
+    LIST_TEXT = "application/list-text"
+    X_TEXT = "application/x-text"
+    TEXT_CSV = "text/csv"
+
+
+class InferenceImageFilename(str, Enum):
+    """Filename of the inference asset in JumpStart distribution buckets."""
+
+    DOG = "dog.jpg"
+    CAT = "cat.jpg"
+    DAISY = "100080576_f52e8ee070_n.jpg"
+    DAISY_2 = "10140303196_b88d3d6cec.jpg"
+    ROSE = "102501987_3cdb8e5394_n.jpg"
+    NAXOS_TAVERNA = "Naxos_Taverna.jpg"
+    PEDESTRIAN = "img_pedestrian.png"
+
+
+class InferenceTabularDataname(str, Enum):
+    """Filename of the tabular data example in JumpStart distribution buckets."""
+
+    REGRESSION_ONEHOT = "regressonehot_data.csv"
+    REGRESSION = "regress_data.csv"
+    MULTICLASS = "multiclass_data.csv"
+
+
+class ClassLabelFile(str, Enum):
+    """Filename in JumpStart distribution buckets for the map of the class index to human readable labels."""
+
+    IMAGE_NET = "ImageNetLabels.txt"
+
+
+TEST_ASSETS_SPECS: Dict[
+    Union[InferenceImageFilename, InferenceTabularDataname, ClassLabelFile], str
+] = {
+    InferenceImageFilename.DOG: _to_s3_path(InferenceImageFilename.DOG, _NB_ASSETS_S3_FOLDER),
+    InferenceImageFilename.CAT: _to_s3_path(InferenceImageFilename.CAT, _NB_ASSETS_S3_FOLDER),
+    InferenceImageFilename.DAISY: _to_s3_path(
+        InferenceImageFilename.DAISY, f"{_TF_FLOWERS_S3_FOLDER}/daisy"
+    ),
+    InferenceImageFilename.DAISY_2: _to_s3_path(
+        InferenceImageFilename.DAISY_2, f"{_TF_FLOWERS_S3_FOLDER}/daisy"
+    ),
+    InferenceImageFilename.ROSE: _to_s3_path(
+        InferenceImageFilename.ROSE, f"{_TF_FLOWERS_S3_FOLDER}/roses"
+    ),
+    InferenceImageFilename.NAXOS_TAVERNA: _to_s3_path(
+        InferenceImageFilename.NAXOS_TAVERNA, _NB_ASSETS_S3_FOLDER
+    ),
+    InferenceImageFilename.PEDESTRIAN: _to_s3_path(
+        InferenceImageFilename.PEDESTRIAN, _NB_ASSETS_S3_FOLDER
+    ),
+    ClassLabelFile.IMAGE_NET: _to_s3_path(ClassLabelFile.IMAGE_NET, _NB_ASSETS_S3_FOLDER),
+    InferenceTabularDataname.REGRESSION_ONEHOT: _to_s3_path(
+        InferenceTabularDataname.REGRESSION_ONEHOT, _NB_ASSETS_S3_FOLDER
+    ),
+    InferenceTabularDataname.REGRESSION: _to_s3_path(
+        InferenceTabularDataname.REGRESSION, _NB_ASSETS_S3_FOLDER
+    ),
+    InferenceTabularDataname.MULTICLASS: _to_s3_path(
+        InferenceTabularDataname.MULTICLASS, _NB_ASSETS_S3_FOLDER
+    ),
+}