change: use recommended inference image URI from Neo API

CHANGELOG.md

@@ -2,9 +2,28 @@
 
 ## v2.72.3 (2022-01-10)
 
+### Features
+
+ * default repack encryption
+ * support large pipeline
+ * add support for pytorch 1.10.0
+
+### Documentation Changes
+
+ * SageMaker model parallel library 1.6.0 API doc
+
 ### Bug Fixes and Other Changes
 
- * update master from dev
+ * Model Registration with BYO scripts
+ * Add ContentType in test_auto_ml_describe
+ * Re-deploy static integ test endpoint if it is not found
+ * fix kmeans test deletion sequence, increment lineage statics
+ * Increment static lineage pipeline
+ * Fix lineage query integ tests
+ * Add label_headers option for Clarify ModelExplainabilityMonitor
+ * Add action type to lineage object
+ * Collapse cross-account artifacts in query lineage response
+ * Update CHANGELOG.md to remove defaulting dot characters
 
 ## v2.72.2 (2022-01-06)
 

setup.py

@@ -104,6 +104,7 @@ def read_version():
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
     ],
     install_requires=required_packages,
     extras_require=extras,

src/sagemaker/image_uris.py

@@ -53,7 +53,7 @@ def retrieve(
         py_version (str): The Python version. This is required if there is
             more than one supported Python version for the given framework version.
         instance_type (str): The SageMaker instance type. For supported types, see
-            https://aws.amazon.com/sagemaker/pricing/instance-types. This is required if
+            https://aws.amazon.com/sagemaker/pricing. This is required if
             there are different images for different processor types.
         accelerator_type (str): Elastic Inference accelerator type. For more, see
             https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.

src/sagemaker/lineage/context.py

@@ -283,6 +283,55 @@ def models(self) -> List[association.Association]:
         ]
         return model_list
 
+    def models_v2(
+        self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.DESCENDANTS
+    ) -> List[Artifact]:
+        """Get artifacts representing models from the context lineage by querying lineage data.
+
+        Args:
+            direction (LineageQueryDirectionEnum, optional): The query direction.
+
+        Returns:
+            list of Artifacts: Artifacts representing a model.
+        """
+        # Firstly query out the model_deployment vertices
+        query_filter = LineageFilter(
+            entities=[LineageEntityEnum.ACTION], sources=[LineageSourceEnum.MODEL_DEPLOYMENT]
+        )
+        model_deployment_query_result = LineageQuery(self.sagemaker_session).query(
+            start_arns=[self.context_arn],
+            query_filter=query_filter,
+            direction=direction,
+            include_edges=False,
+        )
+        if not model_deployment_query_result:
+            return []
+
+        model_deployment_vertices: [] = model_deployment_query_result.vertices
+
+        # Secondary query model based on model deployment
+        model_vertices = []
+        for vertex in model_deployment_vertices:
+            query_result = LineageQuery(self.sagemaker_session).query(
+                start_arns=[vertex.arn],
+                query_filter=LineageFilter(
+                    entities=[LineageEntityEnum.ARTIFACT], sources=[LineageSourceEnum.MODEL]
+                ),
+                direction=LineageQueryDirectionEnum.DESCENDANTS,
+                include_edges=False,
+            )
+            model_vertices.extend(query_result.vertices)
+
+        if not model_vertices:
+            return []
+
+        model_artifacts = []
+        for vertex in model_vertices:
+            lineage_object = vertex.to_lineage_object()
+            model_artifacts.append(lineage_object)
+
+        return model_artifacts
+
     def dataset_artifacts(
         self, direction: LineageQueryDirectionEnum = LineageQueryDirectionEnum.ASCENDANTS
     ) -> List[Artifact]:

src/sagemaker/model.py

@@ -663,12 +663,7 @@ def compile(
             if target_instance_family == "ml_eia2":
                 pass
             elif target_instance_family.startswith("ml_"):
-                self.image_uri = self._compilation_image_uri(
-                    self.sagemaker_session.boto_region_name,
-                    target_instance_family,
-                    framework,
-                    framework_version,
-                )
+                self.image_uri = job_status.get("InferenceImage", None)
                 self._is_compiled_model = True
             else:
                 LOGGER.warning(

src/sagemaker/workflow/lambda_step.py

@@ -154,7 +154,6 @@ def _get_function_arn(self):
         Method creates a lambda function and returns it's arn.
         If the lambda is already present, it will build it's arn and return that.
         """
-        account_id = self.lambda_func.session.account_id()
         region = self.lambda_func.session.boto_region_name
         if region.lower() == "cn-north-1" or region.lower() == "cn-northwest-1":
             partition = "aws-cn"
@@ -163,6 +162,7 @@ def _get_function_arn(self):
 
         if self.lambda_func.function_arn is None:
             try:
+                account_id = self.lambda_func.session.account_id()
                 response = self.lambda_func.create()
                 return response["FunctionArn"]
             except ValueError as error:

tests/integ/sagemaker/lineage/conftest.py

@@ -36,6 +36,7 @@
 from tests.integ.sagemaker.lineage.helpers import name, names
 
 SLEEP_TIME_SECONDS = 1
+SLEEP_TIME_TWO_SECONDS = 2
 STATIC_PIPELINE_NAME = "SdkIntegTestStaticPipeline17"
 STATIC_ENDPOINT_NAME = "SdkIntegTestStaticEndpoint17"
 
@@ -360,12 +361,10 @@ def endpoint_context_obj(sagemaker_session):
 
 @pytest.fixture
 def model_obj(sagemaker_session):
-    model = context.Context.create(
-        context_name=name(),
+    model = artifact.Artifact.create(
+        artifact_name=name(),
+        artifact_type="Model",
         source_uri="bar1",
-        source_type="test-source-type1",
-        context_type="Model",
-        description="test-description",
         properties={"k1": "v1"},
         sagemaker_session=sagemaker_session,
     )
@@ -417,11 +416,12 @@ def endpoint_context_associate_with_model(sagemaker_session, endpoint_action_obj
 
     association.Association.create(
         source_arn=endpoint_action_obj.action_arn,
-        destination_arn=model_obj.context_arn,
+        destination_arn=model_obj.artifact_arn,
         sagemaker_session=sagemaker_session,
     )
     yield obj
-    time.sleep(SLEEP_TIME_SECONDS)
+    # sleep 2 seconds since take longer for lineage injection
+    time.sleep(SLEEP_TIME_TWO_SECONDS)
     obj.delete(disassociate=True)
 
 

tests/integ/sagemaker/lineage/test_endpoint_context.py

@@ -12,17 +12,31 @@
 # language governing permissions and limitations under the License.
 """This module contains code to test SageMaker ``Contexts``"""
 from __future__ import absolute_import
+import time
+
+SLEEP_TIME_ONE_SECONDS = 1
 
 
 def test_model(endpoint_context_associate_with_model, model_obj, endpoint_action_obj):
     model_list = endpoint_context_associate_with_model.models()
     for model in model_list:
         assert model.source_arn == endpoint_action_obj.action_arn
-        assert model.destination_arn == model_obj.context_arn
+        assert model.destination_arn == model_obj.artifact_arn
         assert model.source_type == "ModelDeployment"
         assert model.destination_type == "Model"
 
 
+def test_model_v2(endpoint_context_associate_with_model, model_obj, sagemaker_session):
+    time.sleep(SLEEP_TIME_ONE_SECONDS)
+    model_list = endpoint_context_associate_with_model.models_v2()
+    assert len(model_list) == 1
+    for model in model_list:
+        assert model.artifact_arn == model_obj.artifact_arn
+        assert model.artifact_name == model_obj.artifact_name
+        assert model.artifact_type == "Model"
+        assert model.properties == model_obj.properties
+
+
 def test_dataset_artifacts(static_endpoint_context):
     artifacts_from_query = static_endpoint_context.dataset_artifacts()
 

tests/unit/sagemaker/lineage/test_endpoint_context.py

@@ -15,6 +15,8 @@
 import unittest.mock
 
 from sagemaker.lineage import context, _api_types
+from sagemaker.lineage._api_types import ArtifactSource
+from sagemaker.lineage.artifact import ModelArtifact
 
 
 def test_models(sagemaker_session):
@@ -75,3 +77,113 @@ def test_models(sagemaker_session):
         )
     ]
     assert expected_model_list == model_list
+
+
+def test_models_v2(sagemaker_session):
+    arn1 = "arn:aws:sagemaker:us-west-2:123456789012:context/lineage-integ-3b05f017-0d87-4c37"
+
+    obj = context.EndpointContext(sagemaker_session, context_name="foo", context_arn=arn1)
+
+    sagemaker_session.sagemaker_client.query_lineage.return_value = {
+        "Vertices": [
+            {"Arn": arn1, "Type": "Model", "LineageType": "Artifact"},
+        ],
+        "Edges": [{"SourceArn": "arn1", "DestinationArn": "arn2", "AssociationType": "Produced"}],
+    }
+
+    sagemaker_session.sagemaker_client.describe_context.return_value = {
+        "ContextName": "MyContext",
+        "ContextArn": arn1,
+        "Source": {
+            "SourceUri": "arn:aws:sagemaker:us-west-2:0123456789012:endpoint/myendpoint",
+            "SourceType": "ARN",
+            "SourceId": "Thu Dec 17 17:16:24 UTC 2020",
+        },
+        "ContextType": "Endpoint",
+        "Properties": {
+            "PipelineExecutionArn": "arn:aws:sagemaker:us-west-2:0123456789012:\
+                pipeline/mypipeline/execution/0irnteql64d0",
+            "PipelineStepName": "MyStep",
+            "Status": "Completed",
+        },
+        "CreationTime": 1608225384.0,
+        "CreatedBy": {},
+        "LastModifiedTime": 1608225384.0,
+        "LastModifiedBy": {},
+    }
+
+    sagemaker_session.sagemaker_client.describe_artifact.return_value = {
+        "ArtifactName": "MyArtifact",
+        "ArtifactArn": arn1,
+        "Source": {
+            "SourceUri": "arn:aws:sagemaker:us-west-2:0123456789012:model/mymodel",
+            "SourceType": "ARN",
+            "SourceId": "Thu Dec 17 17:16:24 UTC 2020",
+        },
+        "ArtifactType": "Model",
+        "Properties": {
+            "PipelineExecutionArn": "arn:aws:sagemaker:us-west-2:0123456789012:\
+                pipeline/mypipeline/execution/0irnteql64d0",
+            "PipelineStepName": "MyStep",
+            "Status": "Completed",
+        },
+        "CreationTime": 1608225384.0,
+        "CreatedBy": {},
+        "LastModifiedTime": 1608225384.0,
+        "LastModifiedBy": {},
+    }
+
+    model_list = obj.models_v2()
+
+    expected_calls = [
+        unittest.mock.call(
+            Direction="Descendants",
+            Filters={"Types": ["ModelDeployment"], "LineageTypes": ["Action"]},
+            IncludeEdges=False,
+            MaxDepth=10,
+            StartArns=[arn1],
+        ),
+        unittest.mock.call(
+            Direction="Descendants",
+            Filters={"Types": ["Model"], "LineageTypes": ["Artifact"]},
+            IncludeEdges=False,
+            MaxDepth=10,
+            StartArns=[arn1],
+        ),
+    ]
+    assert expected_calls == sagemaker_session.sagemaker_client.query_lineage.mock_calls
+
+    expected_model_list = [
+        ModelArtifact(
+            artifact_arn=arn1,
+            artifact_name="MyArtifact",
+            source=ArtifactSource(
+                source_uri="arn:aws:sagemaker:us-west-2:0123456789012:model/mymodel",
+                source_types=None,
+                source_type="ARN",
+                source_id="Thu Dec 17 17:16:24 UTC 2020",
+            ),
+            artifact_type="Model",
+            properties={
+                "PipelineExecutionArn": "arn:aws:sagemaker:us-west-2:0123456789012:\
+                pipeline/mypipeline/execution/0irnteql64d0",
+                "PipelineStepName": "MyStep",
+                "Status": "Completed",
+            },
+            creation_time=1608225384.0,
+            created_by={},
+            last_modified_time=1608225384.0,
+            last_modified_by={},
+        )
+    ]
+
+    assert expected_model_list[0].artifact_arn == model_list[0].artifact_arn
+    assert expected_model_list[0].artifact_name == model_list[0].artifact_name
+    assert expected_model_list[0].source == model_list[0].source
+    assert expected_model_list[0].artifact_type == model_list[0].artifact_type
+    assert expected_model_list[0].artifact_type == "Model"
+    assert expected_model_list[0].properties == model_list[0].properties
+    assert expected_model_list[0].creation_time == model_list[0].creation_time
+    assert expected_model_list[0].created_by == model_list[0].created_by
+    assert expected_model_list[0].last_modified_time == model_list[0].last_modified_time
+    assert expected_model_list[0].last_modified_by == model_list[0].last_modified_by