feature: add utils.generate_tensorboard_url function (aws#239)

Added utils.generate_tensorboard_url function to generate TensorBoard URL for given list of s3 bucket paths

Author: pavlo-iuriichuk

src/sagemaker/utils.py

@@ -24,13 +24,18 @@
 import tempfile
 import time
 
+
 from datetime import datetime
 from functools import wraps
 from six.moves.urllib import parse
 
 import six
 
 ECR_URI_PATTERN = r"^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)(amazonaws.com|c2s.ic.gov)(/)(.*:.*)$"
+MAX_BUCKET_PATHS_COUNT = 5
+S3_PREFIX = "s3://"
+HTTP_PREFIX = "http://"
+HTTPS_PREFIX = "https://"
 
 
 # Use the base name of the image as the job name if the user doesn't give us one
@@ -270,6 +275,55 @@ def secondary_training_status_message(job_description, prev_description):
     return "\n".join(status_strs)
 
 
+def generate_tensorboard_url(domain, bucket_paths):
+    """Generate Tensorboard URL for given list of s3 buckets
+
+    Args:
+        domain: JupyterLab app domain
+        bucket_paths: List of S3 bucket paths in format `bucket/path`
+                      or a single string in the same format
+
+    Returns:
+        str: Tensorboard URL
+
+    Raises:
+        AttributeError if invalid inputs are passed
+    """
+
+    def trim_prefix(s, prefix):
+        if s.startswith(prefix):
+            return s[len(prefix) :]
+        return s
+
+    def encode_s3_url(s3_url):
+        if not s3_url:
+            raise AttributeError("bucket_paths element should not be empty")
+        s3_url = trim_prefix(s3_url, S3_PREFIX)
+        return parse.quote_plus("{}{}".format(S3_PREFIX, s3_url))
+
+    if not isinstance(domain, six.string_types):
+        raise AttributeError("domain parameter should be string")
+
+    if len(domain) == 0:
+        raise AttributeError("domain parameter should not be empty")
+
+    if isinstance(bucket_paths, six.string_types):
+        bucket_paths = [bucket_paths]
+    elif not isinstance(bucket_paths, list):
+        raise AttributeError("bucket paths should be a list or a string")
+
+    if len(bucket_paths) == 0:
+        raise AttributeError("bucket_paths parameter should not be empty list")
+
+    domain = trim_prefix(domain, HTTPS_PREFIX)
+    domain = trim_prefix(domain, HTTP_PREFIX)
+
+    s3_urls = map(encode_s3_url, bucket_paths)
+    query = ",".join(s3_urls)
+
+    return "https://{}/tensorboard/default?s3urls={}".format(domain, query)
+
+
 def download_folder(bucket_name, prefix, target, sagemaker_session):
     """Download a folder from S3 to a local path
 

tests/unit/test_utils.py

@@ -208,6 +208,105 @@ def test_secondary_training_status_message_prev_missing():
     )
 
 
+def test_generate_tensorboard_url_valid_domain_and_bucket_paths():
+    domain = "jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = ["bucket1/path1", "bucket2/path2"]
+    expected = "https://{}/tensorboard/default?{}".format(
+        domain, "s3urls=s3%3A%2F%2Fbucket1%2Fpath1,s3%3A%2F%2Fbucket2%2Fpath2"
+    )
+    assert sagemaker.utils.generate_tensorboard_url(domain, bucket_paths) == expected
+
+
+def test_generate_tensorboard_url_valid_domain_and_bucket_paths_with_s3_prefixes():
+    domain = "jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = ["s3://bucket1/path1", "s3://bucket2/path2"]
+    expected = "https://{}/tensorboard/default?{}".format(
+        domain, "s3urls=s3%3A%2F%2Fbucket1%2Fpath1,s3%3A%2F%2Fbucket2%2Fpath2"
+    )
+    assert sagemaker.utils.generate_tensorboard_url(domain, bucket_paths) == expected
+
+
+def test_generate_tensorboard_url_valid_domain_and_bucket_paths_single():
+    domain = "jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = ["bucket1/path1"]
+    expected = "https://{}/tensorboard/default?{}".format(
+        domain, "s3urls=s3%3A%2F%2Fbucket1%2Fpath1"
+    )
+    assert sagemaker.utils.generate_tensorboard_url(domain, bucket_paths) == expected
+
+
+def test_generate_tensorboard_url_valid_domain_and_bucket_paths_string():
+    domain = "jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = "bucket1/path1"
+    expected = "https://{}/tensorboard/default?{}".format(
+        domain, "s3urls=s3%3A%2F%2Fbucket1%2Fpath1"
+    )
+    assert sagemaker.utils.generate_tensorboard_url(domain, bucket_paths) == expected
+
+
+def test_generate_tensorboard_url_valid_domain_with_http_prefix_and_bucket_paths():
+    domain = "http://jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = ["bucket1/path1"]
+    expected = "https://{}/tensorboard/default?{}".format(
+        "jupyterlab.us-east-2.abcdefgh.com", "s3urls=s3%3A%2F%2Fbucket1%2Fpath1"
+    )
+    assert sagemaker.utils.generate_tensorboard_url(domain, bucket_paths) == expected
+
+
+def test_generate_tensorboard_url_valid_domain_with_https_prefix_and_bucket_paths():
+    domain = "https://jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = ["bucket1/path1"]
+    expected = "https://{}/tensorboard/default?{}".format(
+        "jupyterlab.us-east-2.abcdefgh.com", "s3urls=s3%3A%2F%2Fbucket1%2Fpath1"
+    )
+    assert sagemaker.utils.generate_tensorboard_url(domain, bucket_paths) == expected
+
+
+def test_generate_tensorboard_url_bucket_path_neither_string_nor_list():
+    domain = "jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = None
+    try:
+        sagemaker.utils.generate_tensorboard_url(domain, bucket_paths)
+    except AttributeError as error:
+        assert str(error) == "bucket paths should be a list or a string"
+
+
+def test_generate_tensorboard_url_empty_domain():
+    domain = ""
+    bucket_paths = ["bucket1/path1"]
+    try:
+        sagemaker.utils.generate_tensorboard_url(domain, bucket_paths)
+    except AttributeError as error:
+        assert str(error) == "domain parameter should not be empty"
+
+
+def test_generate_tensorboard_url_empty_bucket_paths():
+    domain = "jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = []
+    try:
+        sagemaker.utils.generate_tensorboard_url(domain, bucket_paths)
+    except AttributeError as error:
+        assert str(error) == "bucket_paths parameter should not be empty list"
+
+
+def test_generate_tensorboard_url_bucket_paths_with_empty_string():
+    domain = "jupyterlab.us-east-2.abcdefgh.com"
+    bucket_paths = [""]
+    try:
+        sagemaker.utils.generate_tensorboard_url(domain, bucket_paths)
+    except AttributeError as error:
+        assert str(error) == "bucket_paths element should not be empty"
+
+
+def test_generate_tensorboard_url_domain_non_string():
+    domain = None
+    bucket_paths = ["bucket1/path1"]
+    try:
+        sagemaker.utils.generate_tensorboard_url(domain, bucket_paths)
+    except AttributeError as error:
+        assert str(error) == "domain parameter should be string"
+
+
 @patch("os.makedirs")
 def test_download_folder(makedirs):
     boto_mock = Mock(name="boto_session")