merge integ load test

Author: Yi-Ting Lee

requirements/extras/local_requirements.txt

@@ -1,4 +1,4 @@
 urllib3==1.26.8
 docker-compose==1.29.2
 docker~=5.0.0
-PyYAML==5.4.1
+PyYAML==5.4.1

requirements/extras/test_requirements.txt

@@ -18,3 +18,4 @@ fabric==2.6.0
 requests==2.27.1
 sagemaker-experiments==0.1.35
 Jinja2==3.0.3
+pyvis==0.2.1

src/sagemaker/lineage/query.py

@@ -364,7 +364,7 @@ def _get_visualization_elements(self):
         elements = {"nodes": verts, "edges": edges}
         return elements
 
-    def visualize(self):
+    def visualize(self, path="pyvisExample.html"):
         """Visualize lineage query result."""
         lineage_graph = {
             # nodes can have shape / color
@@ -398,7 +398,7 @@ def visualize(self):
 
         pyvis_vis = PyvisVisualizer(lineage_graph)
         elements = self._get_visualization_elements()
-        return pyvis_vis.render(elements=elements)
+        return pyvis_vis.render(elements=elements, path=path)
 
 
 class LineageFilter(object):

tests/integ/sagemaker/lineage/helpers.py

@@ -12,10 +12,16 @@
 # language governing permissions and limitations under the License.
 """This module contains helper methods for tests of SageMaker Lineage"""
 from __future__ import absolute_import
+from urllib import response
 
 import uuid
 from datetime import datetime
 import time
+import boto3
+from botocore.config import Config
+from sagemaker.lineage import association
+from sagemaker.lineage.artifact import Artifact
+from sagemaker.lineage.association import Association
 
 
 def name():
@@ -78,3 +84,96 @@ def visit(arn, visited: set):
 
     ret = []
     return visit(start_arn, set())
+
+
+class LineageResourceHelper:
+    def __init__(self, sagemaker_session):
+        self.client = sagemaker_session.sagemaker_client
+        self.artifacts = []
+        self.actions = []
+        self.contexts = []
+        self.trialComponents = []
+        self.associations = []
+
+    def create_artifact(self, artifact_name, artifact_type="Dataset"):
+        response = self.client.create_artifact(
+            ArtifactName=artifact_name,
+            Source={
+                "SourceUri": "Test-artifact-" + artifact_name,
+                "SourceTypes": [
+                    {"SourceIdType": "S3ETag", "Value": "Test-artifact-sourceId-value"},
+                ],
+            },
+            ArtifactType=artifact_type,
+        )
+        self.artifacts.append(response["ArtifactArn"])
+
+        return response["ArtifactArn"]
+
+    def create_action(self, action_name, action_type="ModelDeployment"):
+        response = self.client.create_action(
+            ActionName=action_name,
+            Source={
+                "SourceUri": "Test-action-" + action_name,
+                "SourceTypes": [
+                    {"SourceIdType": "S3ETag", "Value": "Test-action-sourceId-value"},
+                ],
+            },
+            ActionType=action_type
+        )
+        self.actions.append(response["ActionArn"])
+
+        return response["ActionArn"]
+
+    def create_context(self, context_name, context_type="Endpoint"):
+        response = self.client.create_context(
+            ContextName=context_name,
+            Source={
+                "SourceUri": "Test-context-" + context_name,
+                "SourceTypes": [
+                    {"SourceIdType": "S3ETag", "Value": "Test-context-sourceId-value"},
+                ],
+            },
+            ContextType=context_type
+        )
+        self.contexts.append(response["ContextArn"])
+
+        return response["ContextArn"]
+
+    def create_trialComponent(self, trialComponent_name, trialComponent_type="TrainingJob"):
+        response = self.client.create_trial_component(
+            TrialComponentName=trialComponent_name,
+            
+        )
+
+    def create_association(self, source_arn, dest_arn, association_type="AssociatedWith"):
+        response = self.client.add_association(
+            SourceArn=source_arn, DestinationArn=dest_arn, AssociationType=association_type
+        )
+        if "SourceArn" in response.keys():
+            self.associations.append((source_arn, dest_arn))
+            return True
+        else:
+            return False
+
+    def clean_all(self):
+        for source, dest in self.associations:
+            try:
+                self.client.delete_association(SourceArn=source, DestinationArn=dest)
+                time.sleep(0.5)
+            except Exception as e:
+                print("skipped " + str(e))
+
+        for artifact_arn in self.artifacts:
+            try:
+                self.client.delete_artifact(ArtifactArn=artifact_arn)
+                time.sleep(0.5)
+            except Exception as e:
+                print("skipped " + str(e))
+
+        for action_arn in self.actions:
+            try:
+                self.client.delete_action(ActionArn=action_arn)
+                time.sleep(0.5)
+            except Exception as e:
+                print("skipped " + str(e))

tests/integ/sagemaker/lineage/test_lineage_visualize.py

@@ -0,0 +1,96 @@
+# 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.
+"""This module contains code to test SageMaker ``LineageQueryResult.visualize()``"""
+
+import datetime
+import logging
+import time
+
+import pytest
+
+import sagemaker.lineage.query
+from sagemaker.lineage.query import LineageQueryDirectionEnum
+from tests.integ.sagemaker.lineage.helpers import name, names, retry, LineageResourceHelper
+
+
+def test_LineageResourceHelper(sagemaker_session):
+    # check if LineageResourceHelper works properly
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
+    try:
+        art1 = lineage_resource_helper.create_artifact(artifact_name=name())
+        art2 = lineage_resource_helper.create_artifact(artifact_name=name())
+        lineage_resource_helper.create_association(source_arn=art1, dest_arn=art2)
+        lineage_resource_helper.clean_all()
+    except Exception as e:
+        print(e)
+        assert False
+
+@pytest.mark.skip("visualizer load test")
+def test_wide_graph_visualize(sagemaker_session):
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
+    wide_graph_root_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+
+    # create wide graph
+    # Artifact ----> Artifact
+    #        \ \ \-> Artifact
+    #         \ \--> Artifact
+    #          \--->  ...
+    try:
+        for i in range(500):
+            artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+            lineage_resource_helper.create_association(source_arn=wide_graph_root_arn, dest_arn=artifact_arn)
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    try:
+        lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
+        lq_result = lq.query(start_arns=[wide_graph_root_arn])
+        lq_result.visualize(path="wideGraph.html")
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    lineage_resource_helper.clean_all()
+
+@pytest.mark.skip("visualizer load test")
+def test_long_graph_visualize(sagemaker_session):
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
+    long_graph_root_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+    last_arn = long_graph_root_arn
+
+    # create long graph
+    # Artifact -> Artifact -> ... -> Artifact
+    try:
+        for i in range(20):
+            new_artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+            lineage_resource_helper.create_association(source_arn=last_arn, dest_arn=new_artifact_arn)
+            last_arn = new_artifact_arn
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    try:
+        lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
+        lq_result = lq.query(start_arns=[long_graph_root_arn], direction=LineageQueryDirectionEnum.DESCENDANTS)
+        # max depth = 10 -> graph rendered only has length of ten (in DESCENDANTS direction)
+        lq_result.visualize(path="longGraph.html")
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    lineage_resource_helper.clean_all()