Merge branch 'master' into mask-creds-local-mode

Author: ahsan-z-khan

doc/frameworks/mxnet/using_mxnet.rst

@@ -377,7 +377,7 @@ It loads the model parameters from a ``model.params`` file in the SageMaker mode
         return net
 
 MXNet on Amazon SageMaker has support for `Elastic Inference <https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html>`__, which allows for inference acceleration to a hosted endpoint for a fraction of the cost of using a full GPU instance.
-In order to load and serve your MXNet model through Amazon Elastic Inference, the MXNet context passed to your MXNet Symbol or Module object within your ``model_fn`` needs to be set to ``eia``, as shown `here <https://docs.aws.amazon.com/dlami/latest/devguide/tutorial-mxnet-elastic-inference.html#ei-mxnet>`__.
+In order to load and serve your MXNet model through Amazon Elastic Inference, import the ``eimx`` Python package and make one change in the code to partition your model and optimize it for the ``EIA`` back end, as shown `here <https://docs.aws.amazon.com/dlami/latest/devguide/tutorial-mxnet-elastic-inference.html#ei-mxnet>`__.
 
 Based on the example above, the following code-snippet shows an example custom ``model_fn`` implementation, which enables loading and serving our MXNet model through Amazon Elastic Inference.
 
@@ -392,11 +392,12 @@ Based on the example above, the following code-snippet shows an example custom `
         Returns:
             mxnet.gluon.nn.Block: a Gluon network (for this example)
         """
-        net = models.get_model('resnet34_v2', ctx=mx.eia(), pretrained=False, classes=10)
-        net.load_params('%s/model.params' % model_dir, ctx=mx.eia())
+        net = models.get_model('resnet34_v2', ctx=mx.cpu(), pretrained=False, classes=10)
+        net.load_params('%s/model.params' % model_dir, ctx=mx.cpu())
+        net.hybridize(backend='EIA', static_alloc=True, static_shape=True)
         return net
 
-The `default_model_fn <https://github.com/aws/sagemaker-mxnet-container/pull/55/files#diff-aabf018d906ed282a3c738377d19a8deR71>`__ loads and serve your model through Elastic Inference, if applicable, within the Amazon SageMaker MXNet containers.
+If you are using MXNet 1.5.1 and earlier, the `default_model_fn <https://github.com/aws/sagemaker-mxnet-container/pull/55/files#diff-aabf018d906ed282a3c738377d19a8deR71>`__ loads and serve your model through Elastic Inference, if applicable, within the Amazon SageMaker MXNet containers.
 
 For more information on how to enable MXNet to interact with Amazon Elastic Inference, see `Use Elastic Inference with MXNet <https://docs.aws.amazon.com/dlami/latest/devguide/tutorial-mxnet-elastic-inference.html>`__.
 

src/sagemaker/lineage/visualizer.py

@@ -105,7 +105,7 @@ def _get_start_arn_from_pipeline_execution_step(self, pipeline_execution_step):
             return None
 
         metadata = pipeline_execution_step["Metadata"]
-        jobs = ["TrainingJob", "ProccessingJob", "TransformJob"]
+        jobs = ["TrainingJob", "ProcessingJob", "TransformJob"]
         for job in jobs:
             if job in metadata and metadata[job]:
                 job_arn = metadata[job]["Arn"]

tests/unit/sagemaker/lineage/conftest.py

@@ -0,0 +1,27 @@
+# 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.lineage import visualizer
+import unittest.mock
+
+
+@pytest.fixture
+def sagemaker_session():
+    return unittest.mock.Mock()
+
+
+@pytest.fixture
+def viz(sagemaker_session):
+    return visualizer.LineageTableVisualizer(sagemaker_session)

tests/unit/sagemaker/lineage/test_action.py

@@ -15,15 +15,9 @@
 import datetime
 import unittest.mock
 
-import pytest
 from sagemaker.lineage import action, _api_types
 
 
-@pytest.fixture
-def sagemaker_session():
-    return unittest.mock.Mock()
-
-
 def test_create(sagemaker_session):
     sagemaker_session.sagemaker_client.create_action.return_value = {
         "ActionArn": "bazz",

tests/unit/sagemaker/lineage/test_artifact.py

@@ -15,15 +15,9 @@
 import datetime
 import unittest.mock
 
-import pytest
 from sagemaker.lineage import artifact, _api_types
 
 
-@pytest.fixture
-def sagemaker_session():
-    return unittest.mock.Mock()
-
-
 def test_create(sagemaker_session):
     sagemaker_session.sagemaker_client.create_artifact.return_value = {
         "ArtifactArn": "bazz",

tests/unit/sagemaker/lineage/test_association.py

@@ -15,15 +15,9 @@
 import datetime
 import unittest.mock
 
-import pytest
 from sagemaker.lineage import association, _api_types
 
 
-@pytest.fixture
-def sagemaker_session():
-    return unittest.mock.Mock()
-
-
 def test_create(sagemaker_session):
     sagemaker_session.sagemaker_client.add_association.return_value = {
         "AssociationArn": "bazz",

tests/unit/sagemaker/lineage/test_endpoint_context.py

@@ -14,15 +14,9 @@
 
 import unittest.mock
 
-import pytest
 from sagemaker.lineage import context, _api_types
 
 
-@pytest.fixture
-def sagemaker_session():
-    return unittest.mock.Mock()
-
-
 def test_models(sagemaker_session):
     obj = context.EndpointContext(sagemaker_session, context_name="foo", context_arn="bazz")
 

tests/unit/sagemaker/lineage/test_model_artifact.py

@@ -14,15 +14,9 @@
 
 import unittest.mock
 
-import pytest
 from sagemaker.lineage import artifact, _api_types
 
 
-@pytest.fixture
-def sagemaker_session():
-    return unittest.mock.Mock()
-
-
 def test_trained_models(sagemaker_session):
     model_artifact_obj = artifact.ModelArtifact(
         sagemaker_session, artifact_arn="model-artifact-arn"

tests/unit/sagemaker/lineage/test_visualizer.py

@@ -1,4 +1,4 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2021 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
@@ -14,33 +14,21 @@
 
 import unittest.mock
 
-import pytest
-from sagemaker.lineage import visualizer
 import pandas as pd
 from collections import OrderedDict
 
 
-@pytest.fixture
-def sagemaker_session():
-    return unittest.mock.Mock()
-
-
-@pytest.fixture
-def vizualizer(sagemaker_session):
-    return visualizer.LineageTableVisualizer(sagemaker_session)
-
-
-def test_friendly_name_short_uri(vizualizer, sagemaker_session):
+def test_friendly_name_short_uri(viz, sagemaker_session):
     uri = "s3://f-069083975568/train.txt"
     arn = "test_arn"
     sagemaker_session.sagemaker_client.describe_artifact.return_value = {
         "Source": {"SourceUri": uri, "SourceTypes": ""}
     }
-    actual_name = vizualizer._get_friendly_name(name=None, arn=arn, entity_type="artifact")
+    actual_name = viz._get_friendly_name(name=None, arn=arn, entity_type="artifact")
     assert uri == actual_name
 
 
-def test_friendly_name_long_uri(vizualizer, sagemaker_session):
+def test_friendly_name_long_uri(viz, sagemaker_session):
     uri = (
         "s3://flintstone-end-to-end-tests-gamma-us-west-2-069083975568/results/canary-auto-1608761252626/"
         "preprocessed-data/tuning_data/train.txt"
@@ -49,12 +37,12 @@ def test_friendly_name_long_uri(vizualizer, sagemaker_session):
     sagemaker_session.sagemaker_client.describe_artifact.return_value = {
         "Source": {"SourceUri": uri, "SourceTypes": ""}
     }
-    actual_name = vizualizer._get_friendly_name(name=None, arn=arn, entity_type="artifact")
+    actual_name = viz._get_friendly_name(name=None, arn=arn, entity_type="artifact")
     expected_name = "s3://.../preprocessed-data/tuning_data/train.txt"
     assert expected_name == actual_name
 
 
-def test_trial_component_name(sagemaker_session, vizualizer):
+def test_trial_component_name(viz, sagemaker_session):
     name = "tc-name"
 
     sagemaker_session.sagemaker_client.describe_trial_component.return_value = {
@@ -90,7 +78,7 @@ def test_trial_component_name(sagemaker_session, vizualizer):
         },
     ]
 
-    df = vizualizer.show(trial_component_name=name)
+    df = viz.show(trial_component_name=name)
 
     sagemaker_session.sagemaker_client.describe_trial_component.assert_called_with(
         TrialComponentName=name,
@@ -121,7 +109,7 @@ def test_trial_component_name(sagemaker_session, vizualizer):
     pd.testing.assert_frame_equal(expected_dataframe, df)
 
 
-def test_model_package_arn(sagemaker_session, vizualizer):
+def test_model_package_arn(viz, sagemaker_session):
     name = "model_package_arn"
 
     sagemaker_session.sagemaker_client.list_artifacts.return_value = {
@@ -157,7 +145,7 @@ def test_model_package_arn(sagemaker_session, vizualizer):
         },
     ]
 
-    df = vizualizer.show(model_package_arn=name)
+    df = viz.show(model_package_arn=name)
 
     sagemaker_session.sagemaker_client.list_artifacts.assert_called_with(
         SourceUri=name,
@@ -188,7 +176,7 @@ def test_model_package_arn(sagemaker_session, vizualizer):
     pd.testing.assert_frame_equal(expected_dataframe, df)
 
 
-def test_endpoint_arn(sagemaker_session, vizualizer):
+def test_endpoint_arn(viz, sagemaker_session):
     name = "endpoint_arn"
 
     sagemaker_session.sagemaker_client.list_contexts.return_value = {
@@ -224,7 +212,7 @@ def test_endpoint_arn(sagemaker_session, vizualizer):
         },
     ]
 
-    df = vizualizer.show(endpoint_arn=name)
+    df = viz.show(endpoint_arn=name)
 
     sagemaker_session.sagemaker_client.list_contexts.assert_called_with(
         SourceUri=name,
@@ -253,3 +241,71 @@ def test_endpoint_arn(sagemaker_session, vizualizer):
     )
 
     pd.testing.assert_frame_equal(expected_dataframe, df)
+
+
+def test_processing_job_pipeline_execution_step(viz, sagemaker_session):
+
+    sagemaker_session.sagemaker_client.list_trial_components.return_value = {
+        "TrialComponentSummaries": [{"TrialComponentArn": "tc-arn"}]
+    }
+
+    sagemaker_session.sagemaker_client.list_associations.side_effect = [
+        {
+            "AssociationSummaries": [
+                {
+                    "SourceArn": "a:b:c:d:e:artifact/src-arn-1",
+                    "SourceName": "source-name-1",
+                    "SourceType": "source-type-1",
+                    "DestinationArn": "a:b:c:d:e:artifact/dest-arn-1",
+                    "DestinationName": "dest-name-1",
+                    "DestinationType": "dest-type-1",
+                    "AssociationType": "type-1",
+                }
+            ]
+        },
+        {
+            "AssociationSummaries": [
+                {
+                    "SourceArn": "a:b:c:d:e:artifact/src-arn-2",
+                    "SourceName": "source-name-2",
+                    "SourceType": "source-type-2",
+                    "DestinationArn": "a:b:c:d:e:artifact/dest-arn-2",
+                    "DestinationName": "dest-name-2",
+                    "DestinationType": "dest-type-2",
+                    "AssociationType": "type-2",
+                }
+            ]
+        },
+    ]
+
+    step = {"Metadata": {"ProcessingJob": {"Arn": "proc-job-arn"}}}
+
+    df = viz.show(pipeline_execution_step=step)
+
+    sagemaker_session.sagemaker_client.list_trial_components.assert_called_with(
+        SourceArn="proc-job-arn",
+    )
+
+    expected_calls = [
+        unittest.mock.call(
+            DestinationArn="tc-arn",
+        ),
+        unittest.mock.call(
+            SourceArn="tc-arn",
+        ),
+    ]
+    assert expected_calls == sagemaker_session.sagemaker_client.list_associations.mock_calls
+
+    expected_dataframe = pd.DataFrame.from_dict(
+        OrderedDict(
+            [
+                ("Name/Source", ["source-name-1", "dest-name-2"]),
+                ("Direction", ["Input", "Output"]),
+                ("Type", ["source-type-1", "dest-type-2"]),
+                ("Association Type", ["type-1", "type-2"]),
+                ("Lineage Type", ["artifact", "artifact"]),
+            ]
+        )
+    )
+
+    pd.testing.assert_frame_equal(expected_dataframe, df)