Merge branch 'master' into master

Author: benieric

src/sagemaker/image_uri_config/sagemaker-base-python.json

@@ -12,6 +12,7 @@
                 "ap-southeast-2": "452832661640",
                 "ap-southeast-3": "276181064229",
                 "ap-southeast-5": "148761635175",
+                "ap-southeast-7": "528757812139",
                 "ca-central-1": "310906938811",
                 "cn-north-1": "390048526115",
                 "cn-northwest-1": "390780980154",
@@ -26,6 +27,7 @@
                 "il-central-1": "380164790875",
                 "me-central-1": "103105715889",
                 "me-south-1": "117516905037",
+                "mx-central-1": "396913743851",
                 "sa-east-1": "782484402741",
                 "us-east-1": "081325390199",
                 "us-east-2": "429704687514",

src/sagemaker/session.py

@@ -4463,6 +4463,49 @@ def wait_for_model_package(self, model_package_name, poll=5):
             )
         return desc
 
+    def get_most_recently_created_approved_model_package(self, model_package_group_name):
+        """Returns the most recently created and Approved model package in a model package group
+
+        Args:
+            model_package_group_name (str): Name or Arn of the model package group
+
+        Returns:
+            dict: Returns a "sagemaker.model.ModelPackage" value.
+        """
+
+        approved_model_packages = self.sagemaker_client.list_model_packages(
+            ModelPackageGroupName=model_package_group_name,
+            ModelApprovalStatus="Approved",
+            SortBy="CreationTime",
+            SortOrder="Descending",
+            MaxResults=1,
+        )
+        next_token = approved_model_packages.get("NextToken")
+
+        while (
+            len(approved_model_packages.get("ModelPackageSummaryList")) == 0
+            and next_token is not None
+            and next_token != ""
+        ):
+            approved_model_packages = self.sagemaker_client.list_model_packages(
+                ModelPackageGroupName=model_package_group_name,
+                ModelApprovalStatus="Approved",
+                SortBy="CreationTime",
+                SortOrder="Descending",
+                MaxResults=1,
+                NextToken=next_token,
+            )
+            next_token = approved_model_packages.get("NextToken")
+
+        if len(approved_model_packages.get("ModelPackageSummaryList")) == 0:
+            return None
+
+        return sagemaker.model.ModelPackage(
+            model_package_arn=approved_model_packages.get("ModelPackageSummaryList")[0].get(
+                "ModelPackageArn"
+            )
+        )
+
     def describe_model(self, name):
         """Calls the DescribeModel API for the given model name.
 

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -170,7 +170,7 @@ def test_jumpstart_gated_model(setup):
 
     model = JumpStartModel(
         model_id=model_id,
-        model_version="3.*",  # version >=3.0.0 stores artifacts in jumpstart-private-cache-* buckets
+        model_version="*",  # version >=3.0.0 stores artifacts in jumpstart-private-cache-* buckets
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
     )
@@ -197,7 +197,7 @@ def test_jumpstart_gated_model_inference_component_enabled(setup):
 
     model = JumpStartModel(
         model_id=model_id,
-        model_version="3.*",  # version >=3.0.0 stores artifacts in jumpstart-private-cache-* buckets
+        model_version="*",  # version >=3.0.0 stores artifacts in jumpstart-private-cache-* buckets
         role=get_sm_session().get_caller_identity_arn(),
         sagemaker_session=get_sm_session(),
     )

tests/integ/test_session.py

@@ -15,7 +15,8 @@
 import boto3
 from botocore.config import Config
 
-from sagemaker import Session
+from sagemaker import Session, ModelPackage
+from sagemaker.utils import unique_name_from_base
 
 CUSTOM_BUCKET_NAME = "this-bucket-should-not-exist"
 
@@ -44,3 +45,62 @@ def test_sagemaker_session_does_not_create_bucket_on_init(
 
     s3 = boto3.resource("s3", region_name=boto_session.region_name)
     assert s3.Bucket(CUSTOM_BUCKET_NAME).creation_date is None
+
+
+def test_sagemaker_session_to_return_most_recent_approved_model_package(sagemaker_session):
+    model_package_group_name = unique_name_from_base("test-model-package-group")
+    approved_model_package = sagemaker_session.get_most_recently_created_approved_model_package(
+        model_package_group_name=model_package_group_name
+    )
+    assert approved_model_package is None
+    sagemaker_session.sagemaker_client.create_model_package_group(
+        ModelPackageGroupName=model_package_group_name
+    )
+    approved_model_package = sagemaker_session.get_most_recently_created_approved_model_package(
+        model_package_group_name=model_package_group_name
+    )
+    assert approved_model_package is None
+    source_uri = "dummy source uri"
+    model_package = sagemaker_session.sagemaker_client.create_model_package(
+        ModelPackageGroupName=model_package_group_name, SourceUri=source_uri
+    )
+    approved_model_package = sagemaker_session.get_most_recently_created_approved_model_package(
+        model_package_group_name=model_package_group_name
+    )
+    assert approved_model_package is None
+    ModelPackage(
+        sagemaker_session=sagemaker_session,
+        model_package_arn=model_package["ModelPackageArn"],
+    ).update_approval_status(approval_status="Approved")
+    approved_model_package = sagemaker_session.get_most_recently_created_approved_model_package(
+        model_package_group_name=model_package_group_name
+    )
+    assert approved_model_package is not None
+    assert approved_model_package.model_package_arn == model_package.get("ModelPackageArn")
+    model_package_2 = sagemaker_session.sagemaker_client.create_model_package(
+        ModelPackageGroupName=model_package_group_name, SourceUri=source_uri
+    )
+    approved_model_package = sagemaker_session.get_most_recently_created_approved_model_package(
+        model_package_group_name=model_package_group_name
+    )
+    assert approved_model_package is not None
+    assert approved_model_package.model_package_arn == model_package.get("ModelPackageArn")
+    ModelPackage(
+        sagemaker_session=sagemaker_session,
+        model_package_arn=model_package_2["ModelPackageArn"],
+    ).update_approval_status(approval_status="Approved")
+    approved_model_package = sagemaker_session.get_most_recently_created_approved_model_package(
+        model_package_group_name=model_package_group_name
+    )
+    assert approved_model_package is not None
+    assert approved_model_package.model_package_arn == model_package_2.get("ModelPackageArn")
+
+    sagemaker_session.sagemaker_client.delete_model_package(
+        ModelPackageName=model_package_2["ModelPackageArn"]
+    )
+    sagemaker_session.sagemaker_client.delete_model_package(
+        ModelPackageName=model_package["ModelPackageArn"]
+    )
+    sagemaker_session.sagemaker_client.delete_model_package_group(
+        ModelPackageGroupName=model_package_group_name
+    )

tests/unit/test_session.py

@@ -7253,3 +7253,35 @@ def test_create_model_package_from_containers_to_create_mpg_if_not_present(sagem
     sagemaker_session.sagemaker_client.create_model_package_group.assert_called_with(
         ModelPackageGroupName="mock-mpg"
     )
+
+
+def test_get_most_recently_created_approved_model_package(sagemaker_session):
+    sagemaker_session.sagemaker_client.list_model_packages.side_effect = [
+        (
+            {
+                "ModelPackageSummaryList": [],
+                "NextToken": "NextToken",
+            }
+        ),
+        (
+            {
+                "ModelPackageSummaryList": [
+                    {
+                        "CreationTime": 1697440162,
+                        "ModelApprovalStatus": "Approved",
+                        "ModelPackageArn": "arn:aws:sagemaker:us-west-2:123456789012:model-package/model-version/3",
+                        "ModelPackageGroupName": "model-version",
+                        "ModelPackageVersion": 3,
+                    },
+                ],
+            }
+        ),
+    ]
+    model_package = sagemaker_session.get_most_recently_created_approved_model_package(
+        model_package_group_name="mpg"
+    )
+    assert model_package is not None
+    assert (
+        model_package.model_package_arn
+        == "arn:aws:sagemaker:us-west-2:123456789012:model-package/model-version/3"
+    )