Merge branch 'dev' into inference-url

Author: mufaddal-rohawala

src/sagemaker/model.py

@@ -178,21 +178,26 @@ def register(
         """
         if self.model_data is None:
             raise ValueError("SageMaker Model Package cannot be created without model data.")
+        if image_uri is not None:
+            self.image_uri = image_uri
+        if model_package_group_name is not None:
+            container_def = self.prepare_container_def()
+        else:
+            container_def = {"Image": self.image_uri, "ModelDataUrl": self.model_data}
 
         model_pkg_args = sagemaker.get_model_package_args(
             content_types,
             response_types,
             inference_instances,
             transform_instances,
-            model_package_name,
-            model_package_group_name,
-            self.model_data,
-            image_uri or self.image_uri,
-            model_metrics,
-            metadata_properties,
-            marketplace_cert,
-            approval_status,
-            description,
+            model_package_name=model_package_name,
+            model_package_group_name=model_package_group_name,
+            model_metrics=model_metrics,
+            metadata_properties=metadata_properties,
+            marketplace_cert=marketplace_cert,
+            approval_status=approval_status,
+            description=description,
+            container_def_list=[container_def],
             drift_check_baselines=drift_check_baselines,
         )
         model_package = self.sagemaker_session.create_model_package_from_containers(

tests/integ/sagemaker/lineage/conftest.py

@@ -36,8 +36,8 @@
 from tests.integ.sagemaker.lineage.helpers import name, names
 
 SLEEP_TIME_SECONDS = 1
-STATIC_PIPELINE_NAME = "SdkIntegTestStaticPipeline14"
-STATIC_ENDPOINT_NAME = "SdkIntegTestStaticEndpoint14"
+STATIC_PIPELINE_NAME = "SdkIntegTestStaticPipeline15"
+STATIC_ENDPOINT_NAME = "SdkIntegTestStaticEndpoint15"
 
 
 @pytest.fixture
@@ -518,6 +518,13 @@ def _get_static_pipeline_execution_arn(sagemaker_session):
 def static_endpoint_context(sagemaker_session, static_pipeline_execution_arn):
     endpoint_arn = get_endpoint_arn_from_static_pipeline(sagemaker_session)
 
+    if endpoint_arn is None:
+        _deploy_static_endpoint(
+            execution_arn=static_pipeline_execution_arn,
+            sagemaker_session=sagemaker_session,
+        )
+        endpoint_arn = get_endpoint_arn_from_static_pipeline(sagemaker_session)
+
     contexts = sagemaker_session.sagemaker_client.list_contexts(SourceUri=endpoint_arn)[
         "ContextSummaries"
     ]
@@ -584,11 +591,17 @@ def static_dataset_artifact(static_model_artifact, sagemaker_session):
 
 
 def get_endpoint_arn_from_static_pipeline(sagemaker_session):
-    endpoint_arn = sagemaker_session.sagemaker_client.describe_endpoint(
-        EndpointName=STATIC_ENDPOINT_NAME
-    )["EndpointArn"]
+    try:
+        endpoint_arn = sagemaker_session.sagemaker_client.describe_endpoint(
+            EndpointName=STATIC_ENDPOINT_NAME
+        )["EndpointArn"]
 
-    return endpoint_arn
+        return endpoint_arn
+    except ClientError as e:
+        error = e.response["Error"]
+        if error["Code"] == "ValidationException":
+            return None
+        raise e
 
 
 def get_model_package_arn_from_static_pipeline(pipeline_execution_arn, sagemaker_session):
@@ -654,7 +667,7 @@ def _deploy_static_endpoint(execution_arn, sagemaker_session):
             sagemaker_session=sagemaker_session,
         )
         model_package.deploy(1, "ml.t2.medium", endpoint_name=STATIC_ENDPOINT_NAME)
-        time.sleep(60)
+        time.sleep(120)
     except ClientError as e:
         if e.response["Error"]["Code"] == "ValidationException":
             print(f"Endpoint {STATIC_ENDPOINT_NAME} already exists. Continuing.")

tests/integ/test_auto_ml.py

@@ -15,12 +15,12 @@
 import os
 
 import pytest
-import tests.integ
-from sagemaker import AutoML, CandidateEstimator, AutoMLInput
-
 from botocore.exceptions import ClientError
+
+import tests.integ
+from sagemaker import AutoML, AutoMLInput, CandidateEstimator
 from sagemaker.utils import unique_name_from_base
-from tests.integ import DATA_DIR, AUTO_ML_DEFAULT_TIMEMOUT_MINUTES, auto_ml_utils
+from tests.integ import AUTO_ML_DEFAULT_TIMEMOUT_MINUTES, DATA_DIR, auto_ml_utils
 from tests.integ.timeout import timeout
 
 ROLE = "SageMakerRole"
@@ -169,6 +169,7 @@ def test_auto_ml_describe_auto_ml_job(sagemaker_session):
                 }
             },
             "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
+            "ContentType": "text/csv;header=present",
         }
     ]
     expected_default_output_config = {
@@ -205,6 +206,7 @@ def test_auto_ml_attach(sagemaker_session):
                 }
             },
             "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
+            "ContentType": "text/csv;header=present",
         }
     ]
     expected_default_output_config = {

tests/integ/test_mxnet.py

@@ -231,6 +231,54 @@ def test_register_model_package(
         sagemaker_session.sagemaker_client.delete_model_package(ModelPackageName=model_package_name)
 
 
+def test_register_model_package_versioned(
+    mxnet_training_job,
+    sagemaker_session,
+    mxnet_inference_latest_version,
+    mxnet_inference_latest_py_version,
+    cpu_instance_type,
+):
+    endpoint_name = "test-mxnet-deploy-model-{}".format(sagemaker_timestamp())
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        desc = sagemaker_session.sagemaker_client.describe_training_job(
+            TrainingJobName=mxnet_training_job
+        )
+        model_package_group_name = "register-model-package-{}".format(sagemaker_timestamp())
+        sagemaker_session.sagemaker_client.create_model_package_group(
+            ModelPackageGroupName=model_package_group_name
+        )
+        model_data = desc["ModelArtifacts"]["S3ModelArtifacts"]
+        script_path = os.path.join(DATA_DIR, "mxnet_mnist", "mnist.py")
+        model = MXNetModel(
+            model_data,
+            "SageMakerRole",
+            entry_point=script_path,
+            py_version=mxnet_inference_latest_py_version,
+            sagemaker_session=sagemaker_session,
+            framework_version=mxnet_inference_latest_version,
+        )
+        model_pkg = model.register(
+            content_types=["application/json"],
+            response_types=["application/json"],
+            inference_instances=["ml.m5.large"],
+            transform_instances=["ml.m5.large"],
+            model_package_group_name=model_package_group_name,
+            approval_status="Approved",
+        )
+        assert isinstance(model_pkg, ModelPackage)
+        predictor = model.deploy(1, cpu_instance_type, endpoint_name=endpoint_name)
+        data = numpy.zeros(shape=(1, 1, 28, 28))
+        result = predictor.predict(data)
+        assert result is not None
+        sagemaker_session.sagemaker_client.delete_model_package(
+            ModelPackageName=model_pkg.model_package_arn
+        )
+        sagemaker_session.sagemaker_client.delete_model_package_group(
+            ModelPackageGroupName=model_package_group_name
+        )
+
+
 def test_deploy_model_with_tags_and_kms(
     mxnet_training_job,
     sagemaker_session,