friendly names for short URIs

src/sagemaker/lineage/visualizer.py

@@ -13,8 +13,6 @@
 """This module contains functionality to display lineage data."""
 from __future__ import absolute_import
 import logging
-import os
-from urllib.parse import urlparse
 import pandas as pd
 
 from sagemaker.lineage.association import Association
@@ -303,13 +301,12 @@ def _get_friendly_name(self, name, arn, entity_type):
         if entity_type == "artifact":
             artifact = self._session.sagemaker_client.describe_artifact(ArtifactArn=arn)
             uri = artifact["Source"]["SourceUri"]
-            # try to get file name from url
-            uri_parsed = urlparse(uri)
-            name = os.path.basename(uri_parsed.path)
 
-            # directory?
-            ext = os.path.splitext(name)[1]
-            if not ext or len(ext) > 3:
+            # shorten the uri if the length is more than 40,
+            # e.g s3://flintstone-end-to-end-tests-gamma-us-west-2-069083975568/results/
+            # canary-auto-1608761252626/preprocessed-data/tuning_data/train.txt
+            # become s3://.../preprocessed-data/tuning_data/train.txt
+            if len(uri) > 48:
                 name = uri[:5] + "..." + uri[len(uri) - 40 :]
 
             # if not then use the full uri

tests/unit/sagemaker/lineage/test_visualizer.py

@@ -30,6 +30,30 @@ def vizualizer(sagemaker_session):
     return visualizer.LineageTableVisualizer(sagemaker_session)
 
 
+def test_friendly_name_short_uri(vizualizer, 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")
+    assert uri == actual_name
+
+
+def test_friendly_name_long_uri(vizualizer, sagemaker_session):
+    uri = (
+        "s3://flintstone-end-to-end-tests-gamma-us-west-2-069083975568/results/canary-auto-1608761252626/"
+        "preprocessed-data/tuning_data/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")
+    expected_name = "s3://.../preprocessed-data/tuning_data/train.txt"
+    assert expected_name == actual_name
+
+
 def test_trial_component_name(sagemaker_session, vizualizer):
     name = "tc-name"