Fix tests and codestyle (#1619)

Co-authored-by: Erick Benitez-Ramos <[email protected]>
Co-authored-by: Erick Benitez-Ramos <[email protected]>

Author: 2 people

src/sagemaker/modules/train/model_trainer.py

@@ -893,7 +893,7 @@ def from_recipe(
             sagemaker_session = Session()
             logger.warning("SageMaker session not provided. Using default Session.")
         if role is None:
-            role = get_execution_role(sagemaker_session=session)
+            role = get_execution_role(sagemaker_session=sagemaker_session)
             logger.warning(f"Role not provided. Using default role:\n{role}")
 
         # The training recipe is used to prepare the following args:

src/sagemaker/serve/builder/model_builder.py

@@ -838,14 +838,17 @@ def _initialize_for_mlflow(self, artifact_path: str) -> None:
 
     @_capture_telemetry("ModelBuilder.build_training_job")
     def _collect_training_job_model_telemetry(self):
+        """Dummy method to collect telemetry for training job handshake"""
         return
 
     @_capture_telemetry("ModelBuilder.build_model_trainer")
     def _collect_model_trainer_model_telemetry(self):
+        """Dummy method to collect telemetry for model trainer handshake"""
         return
 
     @_capture_telemetry("ModelBuilder.build_estimator")
     def _collect_estimator_model_telemetry(self):
+        """Dummy method to collect telemetry for estimator handshake"""
         return
 
     # Model Builder is a class to build the model for deployment.

src/sagemaker/serve/utils/telemetry_logger.py

@@ -122,13 +122,12 @@ def wrapper(self, *args, **kwargs):
                 extra += f"&x-modelServer={MODEL_SERVER_TO_CODE[str(self.model_server)]}"
 
             if self.image_uri:
-                image_uri_tail = self.image_uri.split("/")[1]
                 image_uri_option = _get_image_uri_option(
                     self.image_uri, getattr(self, "_is_custom_image_uri", False)
                 )
-
-            if self.image_uri:
-                extra += f"&x-imageTag={image_uri_tail}"
+                split_image_uri = self.image_uri.split("/")
+                if len(split_image_uri) > 1:
+                    extra += f"&x-imageTag={split_image_uri[1]}"
 
             extra += f"&x-sdkVersion={SDK_VERSION}"
 

src/sagemaker/telemetry/telemetry_logging.py

@@ -63,7 +63,8 @@
 
 
 def _telemetry_emitter(feature: str, func_name: str):
-    """
+    """Telemetry Emitter
+
     Decorator to emit telemetry logs for SageMaker Python SDK functions. This class needs
     sagemaker_session object as a member. Default session object is a pysdk v2 Session object
     in this repo. When collecting telemetry for classes using sagemaker-core Session object,

tests/integ/sagemaker/modules/conftest.py

@@ -36,5 +36,5 @@ def modules_sagemaker_session():
 
     yield sagemaker_session
 
-    if region_manual_set:
+    if region_manual_set and "AWS_DEFAULT_REGION" in os.environ:
         del os.environ["AWS_DEFAULT_REGION"]

tests/integ/sagemaker/modules/train/test_model_trainer.py

@@ -44,7 +44,7 @@ def test_hp_contract_basic_py_script(modules_sagemaker_session):
     )
 
     model_trainer = ModelTrainer(
-        session=modules_sagemaker_session,
+        sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
         hyperparameters=EXPECTED_HYPERPARAMETERS,
         source_code=source_code,
@@ -60,7 +60,7 @@ def test_hp_contract_basic_sh_script(modules_sagemaker_session):
         entry_script="train.sh",
     )
     model_trainer = ModelTrainer(
-        session=modules_sagemaker_session,
+        sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
         hyperparameters=EXPECTED_HYPERPARAMETERS,
         source_code=source_code,
@@ -77,7 +77,7 @@ def test_hp_contract_mpi_script(modules_sagemaker_session):
     )
     compute = Compute(instance_type="ml.m5.xlarge", instance_count=2)
     model_trainer = ModelTrainer(
-        session=modules_sagemaker_session,
+        sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
         compute=compute,
         hyperparameters=EXPECTED_HYPERPARAMETERS,
@@ -96,7 +96,7 @@ def test_hp_contract_torchrun_script(modules_sagemaker_session):
     )
     compute = Compute(instance_type="ml.m5.xlarge", instance_count=2)
     model_trainer = ModelTrainer(
-        session=modules_sagemaker_session,
+        sagemaker_session=modules_sagemaker_session,
         training_image=DEFAULT_CPU_IMAGE,
         compute=compute,
         hyperparameters=EXPECTED_HYPERPARAMETERS,

tests/integ/sagemaker/serve/conftest.py

@@ -10,64 +10,48 @@
 # 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 __future__ import absolute_import
 
-# import os
-# import pytest
-# import platform
-# import collections
-# from numpy import loadtxt
-# from sagemaker.serve.spec.inference_spec import InferenceSpec
+import pytest
+import os
+import boto3
+import sagemaker
+import sagemaker_core.helper.session_helper as core_session
 
-# if platform.python_version_tuple()[1] == "8":
-#     from xgboost import XGBClassifier
-#     from sklearn.model_selection import train_test_split
+DEFAULT_REGION = "us-west-2"
 
-# from tests.integ.sagemaker.serve.constants import XGB_RESOURCE_DIR
 
+@pytest.fixture(scope="module")
+def mb_sagemaker_session():
+    region = os.environ.get("AWS_DEFAULT_REGION")
+    if not region:
+        os.environ["AWS_DEFAULT_REGION"] = DEFAULT_REGION
+        region_manual_set = True
+    else:
+        region_manual_set = True
 
-# XgbTestSplit = collections.namedtuple("XgbTrainTestSplit", "x_test y_test")
+    boto_session = boto3.Session(region_name=os.environ["AWS_DEFAULT_REGION"])
+    sagemaker_session = sagemaker.Session(boto_session=boto_session)
 
+    yield sagemaker_session
 
-# @pytest.fixture(scope="session")
-# def loaded_xgb_model():
-#     model = XGBClassifier()
-#     model.load_model(XGB_RESOURCE_DIR + "/model.xgb")
-#     return model
+    if region_manual_set and "AWS_DEFAULT_REGION" in os.environ:
+        del os.environ["AWS_DEFAULT_REGION"]
 
 
-# @pytest.fixture(scope="session")
-# def xgb_inference_spec():
-#     class MyXGBoostModel(InferenceSpec):
-#         def load(self, model_dir: str):
-#             model = XGBClassifier()
-#             model.load_model(model_dir + "/model.xgb")
-#             return model
+@pytest.fixture(scope="module")
+def mb_sagemaker_core_session():
+    region = os.environ.get("AWS_DEFAULT_REGION")
+    if not region:
+        os.environ["AWS_DEFAULT_REGION"] = DEFAULT_REGION
+        region_manual_set = True
+    else:
+        region_manual_set = True
 
-#         def invoke(
-#             self,
-#             input: object,
-#             model: object,
-#         ):
-#             y_pred = model.predict(input)
-#             predictions = [round(value) for value in y_pred]
-#             return predictions
+    boto_session = boto3.Session(region_name=os.environ["AWS_DEFAULT_REGION"])
+    sagemaker_session = core_session.Session(boto_session=boto_session)
 
-#     return MyXGBoostModel()
+    yield sagemaker_session
 
-
-# @pytest.fixture(scope="session")
-# def xgb_test_sets():
-#     dataset = loadtxt(
-#         os.path.join(XGB_RESOURCE_DIR, "classification_training_data.data.csv"), delimiter=","
-#     )
-
-#     X = dataset[:, 0:8]
-#     Y = dataset[:, 8]
-
-#     seed = 7
-#     test_size = 0.33
-
-#     _, x_test, _, y_test = train_test_split(X, Y, test_size=test_size, random_state=seed)
-
-#     return XgbTestSplit(x_test, y_test)
+    if region_manual_set and "AWS_DEFAULT_REGION" in os.environ:
+        del os.environ["AWS_DEFAULT_REGION"]

tests/integ/sagemaker/serve/test_base_model_builder_deploy.py

@@ -0,0 +1,193 @@
+# 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 import get_execution_role
+from sklearn.datasets import load_iris
+from sklearn.model_selection import train_test_split
+
+import os
+
+from sagemaker_core.main.shapes import (
+    AlgorithmSpecification,
+    Channel,
+    DataSource,
+    S3DataSource,
+    OutputDataConfig,
+    ResourceConfig,
+    StoppingCondition,
+)
+import uuid
+from sagemaker.serve.builder.model_builder import ModelBuilder
+import pandas as pd
+import numpy as np
+from sagemaker.serve import InferenceSpec, SchemaBuilder
+from sagemaker_core.main.resources import TrainingJob
+from xgboost import XGBClassifier
+
+from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
+
+from sagemaker.s3_utils import s3_path_join
+from sagemaker.async_inference import AsyncInferenceConfig
+from tests.integ.utils import cleanup_model_resources
+
+
+@pytest.fixture(scope="module")
+def xgboost_model_builder(mb_sagemaker_session):
+    sagemaker_session = mb_sagemaker_session
+    role = get_execution_role(sagemaker_session=sagemaker_session)
+    bucket = sagemaker_session.default_bucket()
+
+    # Get IRIS Data
+    iris = load_iris()
+    iris_df = pd.DataFrame(iris.data, columns=iris.feature_names)
+    iris_df["target"] = iris.target
+
+    # Prepare Data
+    os.makedirs("data", exist_ok=True)
+
+    iris_df = iris_df[["target"] + [col for col in iris_df.columns if col != "target"]]
+
+    train_data, test_data = train_test_split(iris_df, test_size=0.2, random_state=42)
+
+    train_data.to_csv("data/train.csv", index=False, header=False)
+    test_data.to_csv("data/test.csv", index=False, header=False)
+
+    # Remove the target column from the testing data. We will use this to call invoke_endpoint later
+    test_data.drop("target", axis=1)
+
+    prefix = "DEMO-scikit-iris"
+    TRAIN_DATA = "train.csv"
+    DATA_DIRECTORY = "data"
+
+    sagemaker_session.upload_data(
+        DATA_DIRECTORY, bucket=bucket, key_prefix="{}/{}".format(prefix, DATA_DIRECTORY)
+    )
+
+    s3_input_path = "s3://{}/{}/data/{}".format(bucket, prefix, TRAIN_DATA)
+    s3_output_path = "s3://{}/{}/output".format(bucket, prefix)
+
+    print(s3_input_path)
+    print(s3_output_path)
+
+    image = "433757028032.dkr.ecr.us-west-2.amazonaws.com/xgboost:1"
+
+    class XGBoostSpec(InferenceSpec):
+        def load(self, model_dir: str):
+            print(model_dir)
+            model = XGBClassifier()
+            model.load_model(model_dir + "/xgboost-model")
+            return model
+
+        def invoke(self, input_object: object, model: object):
+            prediction_probabilities = model.predict_proba(input_object)
+            predictions = np.argmax(prediction_probabilities, axis=1)
+            return predictions
+
+    data = {"Name": ["Alice", "Bob", "Charlie"]}
+    df = pd.DataFrame(data)
+    training_job_name = str(uuid.uuid4())
+    schema_builder = SchemaBuilder(sample_input=df, sample_output=df)
+
+    training_job = TrainingJob.create(
+        training_job_name=training_job_name,
+        hyper_parameters={
+            "objective": "multi:softmax",
+            "num_class": "3",
+            "num_round": "10",
+            "eval_metric": "merror",
+        },
+        algorithm_specification=AlgorithmSpecification(
+            training_image=image, training_input_mode="File"
+        ),
+        role_arn=role,
+        input_data_config=[
+            Channel(
+                channel_name="train",
+                content_type="csv",
+                compression_type="None",
+                record_wrapper_type="None",
+                data_source=DataSource(
+                    s3_data_source=S3DataSource(
+                        s3_data_type="S3Prefix",
+                        s3_uri=s3_input_path,
+                        s3_data_distribution_type="FullyReplicated",
+                    )
+                ),
+            )
+        ],
+        output_data_config=OutputDataConfig(s3_output_path=s3_output_path),
+        resource_config=ResourceConfig(
+            instance_type="ml.m4.xlarge", instance_count=1, volume_size_in_gb=30
+        ),
+        stopping_condition=StoppingCondition(max_runtime_in_seconds=600),
+    )
+    training_job.wait()
+
+    xgboost_model_builder = ModelBuilder(
+        name="ModelBuilderTest",
+        model_path=training_job.model_artifacts.s3_model_artifacts,
+        role_arn=role,
+        inference_spec=XGBoostSpec(),
+        image_uri=image,
+        schema_builder=schema_builder,
+        instance_type="ml.c6i.xlarge",
+    )
+    xgboost_model_builder.build()
+    return xgboost_model_builder
+
+
+def test_real_time_deployment(xgboost_model_builder):
+    real_time_predictor = xgboost_model_builder.deploy(
+        endpoint_name="test", initial_instance_count=1
+    )
+
+    assert real_time_predictor is not None
+    cleanup_model_resources(
+        sagemaker_session=xgboost_model_builder.sagemaker_session,
+        model_name=xgboost_model_builder.built_model.name,
+        endpoint_name=xgboost_model_builder.built_model.endpoint_name,
+    )
+
+
+def test_serverless_deployment(xgboost_model_builder):
+    serverless_predictor = xgboost_model_builder.deploy(
+        endpoint_name="test1", inference_config=ServerlessInferenceConfig()
+    )
+
+    assert serverless_predictor is not None
+    cleanup_model_resources(
+        sagemaker_session=xgboost_model_builder.sagemaker_session,
+        model_name=xgboost_model_builder.built_model.name,
+        endpoint_name=xgboost_model_builder.built_model.endpoint_name,
+    )
+
+
+def test_async_deployment(xgboost_model_builder, mb_sagemaker_session):
+    async_predictor = xgboost_model_builder.deploy(
+        endpoint_name="test2",
+        inference_config=AsyncInferenceConfig(
+            output_path=s3_path_join(
+                "s3://", mb_sagemaker_session.default_bucket(), "async_inference/output"
+            )
+        ),
+    )
+
+    assert async_predictor is not None
+    cleanup_model_resources(
+        sagemaker_session=xgboost_model_builder.sagemaker_session,
+        model_name=xgboost_model_builder.built_model.name,
+        endpoint_name=xgboost_model_builder.built_model.endpoint_name,
+    )