Merge branch 'master' into sagemaker-mlflow-extras

Author: ryansteakley

.github/PULL_REQUEST_TEMPLATE.md

@@ -22,5 +22,6 @@ _Put an `x` in the boxes that apply. You can also fill these out after creating
 - [ ] I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes
 - [ ] I have checked that my tests are not configured for a specific region or account (if appropriate)
 - [ ] I have used [`unique_name_from_base`](https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/utils.py#L77) to create resource names in integ tests (if appropriate)
+- [ ] If adding any dependency in requirements.txt files, I have spell checked and ensured they exist in PyPi
 
 By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

src/sagemaker/pipeline.py

@@ -26,6 +26,7 @@
 )
 from sagemaker.drift_check_baselines import DriftCheckBaselines
 from sagemaker.metadata_properties import MetadataProperties
+from sagemaker.model import ModelPackage
 from sagemaker.model_card import (
     ModelCard,
     ModelPackageModelCard,
@@ -470,7 +471,18 @@ def register(
             model_card=model_card,
         )
 
-        self.sagemaker_session.create_model_package_from_containers(**model_pkg_args)
+        model_package = self.sagemaker_session.create_model_package_from_containers(
+            **model_pkg_args
+        )
+
+        if model_package is not None and "ModelPackageArn" in model_package:
+            return ModelPackage(
+                role=self.role,
+                model_package_arn=model_package.get("ModelPackageArn"),
+                sagemaker_session=self.sagemaker_session,
+                predictor_cls=self.predictor_cls,
+            )
+        return None
 
     def transformer(
         self,

src/sagemaker/serve/utils/conda_in_process.yml

@@ -37,7 +37,7 @@ dependencies:
             - botocore>=1.29.114
             - cachetools>=5.3.0
             - certifi==2022.12.7
-            - harset-normalizer>=3.1.0
+            - charset-normalizer>=3.1.0
             - click>=8.1.3
             - cloudpickle>=2.2.1
             - colorama>=0.4.4

src/sagemaker/serve/utils/in_process_requirements.txt

@@ -5,7 +5,7 @@ blinker>=1.6.2
 botocore>=1.29.114
 cachetools>=5.3.0
 certifi==2024.7.4
-harset-normalizer>=3.1.0
+charset-normalizer>=3.1.0
 click>=8.1.3
 cloudpickle>=2.2.1
 colorama>=0.4.4

tests/integ/test_inference_pipeline.py

@@ -150,6 +150,40 @@ def test_inference_pipeline_model_deploy(sagemaker_session, cpu_instance_type):
         assert "Could not find model" in str(exception.value)
 
 
+@pytest.mark.release
+def test_inference_pipeline_model_register(sagemaker_session):
+    sparkml_data_path = os.path.join(DATA_DIR, "sparkml_model")
+    endpoint_name = unique_name_from_base("test-inference-pipeline-deploy")
+    sparkml_model_data = sagemaker_session.upload_data(
+        path=os.path.join(sparkml_data_path, "mleap_model.tar.gz"),
+        key_prefix="integ-test-data/sparkml/model",
+    )
+
+    sparkml_model = SparkMLModel(
+        model_data=sparkml_model_data,
+        env={"SAGEMAKER_SPARKML_SCHEMA": SCHEMA},
+        sagemaker_session=sagemaker_session,
+    )
+
+    model = PipelineModel(
+        models=[sparkml_model],
+        role="SageMakerRole",
+        sagemaker_session=sagemaker_session,
+        name=endpoint_name,
+    )
+    model_package_group_name = unique_name_from_base("pipeline-model-package")
+    model_package = model.register(model_package_group_name=model_package_group_name)
+    assert model_package.model_package_arn is not None
+
+    sagemaker_session.sagemaker_client.delete_model_package(
+        ModelPackageName=model_package.model_package_arn
+    )
+
+    sagemaker_session.sagemaker_client.delete_model_package_group(
+        ModelPackageGroupName=model_package_group_name
+    )
+
+
 @pytest.mark.slow_test
 @pytest.mark.flaky(reruns=5, reruns_delay=2)
 def test_inference_pipeline_model_deploy_and_update_endpoint(

tests/unit/sagemaker/image_uris/test_graviton.py

@@ -30,11 +30,18 @@
 ]
 
 
-def _test_graviton_framework_uris(framework, version, py_version, account, region):
+def _test_graviton_framework_uris(
+    framework, version, py_version, account, region, container_version="ubuntu20.04-sagemaker"
+):
     for instance_type in GRAVITON_INSTANCE_TYPES:
         uri = image_uris.retrieve(framework, region, instance_type=instance_type, version=version)
         expected = _expected_graviton_framework_uri(
-            framework, version, py_version, account, region=region
+            framework,
+            version,
+            py_version,
+            account,
+            region=region,
+            container_version=container_version,
         )
         assert expected == uri
 
@@ -50,11 +57,21 @@ def test_graviton_framework_uris(load_config_and_file_name, scope):
     for version in VERSIONS:
         ACCOUNTS = config[scope]["versions"][version]["registries"]
         py_versions = config[scope]["versions"][version]["py_versions"]
+        container_version = (
+            config[scope]["versions"][version].get("container_version", {}).get("cpu", None)
+        )
+        if container_version:
+            container_version = container_version + "-sagemaker"
         for py_version in py_versions:
             for region in ACCOUNTS.keys():
-                _test_graviton_framework_uris(
-                    framework, version, py_version, ACCOUNTS[region], region
-                )
+                if container_version:
+                    _test_graviton_framework_uris(
+                        framework, version, py_version, ACCOUNTS[region], region, container_version
+                    )
+                else:
+                    _test_graviton_framework_uris(
+                        framework, version, py_version, ACCOUNTS[region], region
+                    )
 
 
 def _test_graviton_unsupported_framework(framework, region, framework_version):
@@ -183,11 +200,14 @@ def test_graviton_sklearn_image_scope_specified_x86_instance(graviton_sklearn_un
             assert "Unsupported instance type: m5." in str(error)
 
 
-def _expected_graviton_framework_uri(framework, version, py_version, account, region):
+def _expected_graviton_framework_uri(
+    framework, version, py_version, account, region, container_version
+):
     return expected_uris.graviton_framework_uri(
         "{}-inference-graviton".format(framework),
         fw_version=version,
         py_version=py_version,
         account=account,
         region=region,
+        container_version=container_version,
     )

tests/unit/test_pipeline_model.py

@@ -420,3 +420,27 @@ def test_network_isolation(tfo, time, sagemaker_session):
         vpc_config=None,
         enable_network_isolation=True,
     )
+
+
+def test_pipeline_model_register(sagemaker_session):
+    sagemaker_session.create_model_package_from_containers = Mock(
+        name="create_model_package_from_containers",
+        return_value={
+            "ModelPackageArn": "arn:aws:sagemaker:us-west-2:123456789123:model-package/unit-test-package-version/1"
+        },
+    )
+    framework_model = DummyFrameworkModel(sagemaker_session)
+    sparkml_model = SparkMLModel(
+        model_data=MODEL_DATA_2, role=ROLE, sagemaker_session=sagemaker_session
+    )
+    model = PipelineModel(
+        models=[framework_model, sparkml_model],
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        enable_network_isolation=True,
+    )
+    model_package = model.register()
+    assert (
+        model_package.model_package_arn
+        == "arn:aws:sagemaker:us-west-2:123456789123:model-package/unit-test-package-version/1"
+    )