feature: Introduce TensorBoard app class (#3810)

waytrue17 · Wei Chu · web-flow · commit 98134d23f81c · 2023-04-25T16:28:39.000-04:00
Co-authored-by: Wei Chu &lt;weichu@amazon.com&gt;
diff --git a/src/sagemaker/interactive_apps/tensorboard.py b/src/sagemaker/interactive_apps/tensorboard.py
@@ -0,0 +1,141 @@
+# 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 methods for starting up and accessing TensorBoard apps hosted on SageMaker"""
+from __future__ import absolute_import
+
+import json
+import logging
+import os
+import re
+
+from typing import Optional
+from sagemaker.session import Session, NOTEBOOK_METADATA_FILE
+
+logger = logging.getLogger(__name__)
+
+
+class TensorBoardApp(object):
+    """TensorBoardApp is a class for creating/accessing a TensorBoard app hosted on SageMaker."""
+
+    def __init__(self, region: Optional[str] = None):
+        """Initialize a TensorBoardApp object.
+
+        Args:
+            region (str): The AWS Region, e.g. us-east-1. If not specified,
+                one is created using the default AWS configuration chain.
+        """
+        if region:
+            self.region = region
+        else:
+            try:
+                self.region = Session().boto_region_name
+            except ValueError:
+                raise ValueError(
+                    "Failed to get the Region information from the default config. Please either "
+                    "pass your Region manually as an input argument or set up the local AWS configuration."
+                )
+
+        self._domain_id = None
+        self._user_profile_name = None
+        self._valid_domain_and_user = False
+        self._get_domain_and_user()
+
+    def __str__(self):
+        """Return str(self)."""
+        return f"TensorBoardApp(region={self.region})"
+
+    def __repr__(self):
+        """Return repr(self)."""
+        return self.__str__()
+
+    def get_app_url(self, training_job_name: Optional[str] = None):
+        """Generates an unsigned URL to help access the TensorBoard application hosted in SageMaker.
+
+           For users that are already in SageMaker Studio, this method tries to get the domain id and the user
+           profile from the Studio environment. If succeeded, the generated URL will direct to the TensorBoard
+           application in SageMaker. Otherwise, it will direct to the TensorBoard landing page in the SageMaker
+           console. For non-Studio users, the URL will direct to the TensorBoard landing page in the SageMaker
+           console.
+
+        Args:
+            training_job_name (str): Optional. The name of the training job to pre-load in TensorBoard.
+                If nothing provided, the method still returns the TensorBoard application URL,
+                but the application will not have any training jobs added for tracking. You can
+                add training jobs later by using the SageMaker Data Manager UI.
+                Default: ``None``
+
+        Returns:
+            str: An unsigned URL for TensorBoard hosted on SageMaker.
+        """
+        if self._valid_domain_and_user:
+            url = "https://{}.studio.{}.sagemaker.aws/tensorboard/default".format(
+                self._domain_id, self.region
+            )
+            if training_job_name is not None:
+                self._validate_job_name(training_job_name)
+                url += "/data/plugin/sagemaker_data_manager/add_folder_or_job?Redirect=True&Name={}".format(
+                    training_job_name
+                )
+            else:
+                url += "/#sagemaker_data_manager"
+        else:
+            url = "https://{region}.console.aws.amazon.com/sagemaker/home?region={region}#/tensor-board-landing".format(
+                region=self.region
+            )
+            if training_job_name is not None:
+                self._validate_job_name(training_job_name)
+                url += "/{}".format(training_job_name)
+
+        return url
+
+    def _get_domain_and_user(self):
+        """Get and validate studio domain id and user profile from NOTEBOOK_METADATA_FILE in studio environment.
+
+        Set _valid_domain_and_user to True if validation succeeded.
+        """
+        if not os.path.isfile(NOTEBOOK_METADATA_FILE):
+            return
+
+        with open(NOTEBOOK_METADATA_FILE, "rb") as f:
+            metadata = json.loads(f.read())
+            self._domain_id = metadata.get("DomainId")
+            self._user_profile_name = metadata.get("UserProfileName")
+            if self._validate_domain_id() is True and self._validate_user_profile_name() is True:
+                self._valid_domain_and_user = True
+            else:
+                logger.warning(
+                    "NOTEBOOK_METADATA_FILE detected but failed to get valid domain and user from it."
+                )
+
+    def _validate_job_name(self, job_name: str):
+        """Validate training job name format."""
+        job_name_regex = "^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,62}"
+        if not re.fullmatch(job_name_regex, job_name):
+            raise ValueError(
+                "Invalid job name. Job name must match regular expression {}".format(job_name_regex)
+            )
+
+    def _validate_domain_id(self):
+        """Validate domain id format."""
+        if self._domain_id is None or len(self._domain_id) > 63:
+            return False
+        return True
+
+    def _validate_user_profile_name(self):
+        """Validate user profile name format."""
+        user_profile_name_regex = "^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,62}"
+        if self._user_profile_name is None or not re.fullmatch(
+            user_profile_name_regex, self._user_profile_name
+        ):
+            return False
+        return True
diff --git a/tests/unit/test_tensorboard.py b/tests/unit/test_tensorboard.py
@@ -0,0 +1,139 @@
+# 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
+
+from sagemaker.interactive_apps.tensorboard import TensorBoardApp
+from unittest.mock import patch, mock_open, PropertyMock
+
+import json
+import pytest
+
+TEST_DOMAIN = "testdomain"
+TEST_USER_PROFILE = "testuser"
+TEST_REGION = "testregion"
+TEST_NOTEBOOK_METADATA = json.dumps({"DomainId": TEST_DOMAIN, "UserProfileName": TEST_USER_PROFILE})
+TEST_TRAINING_JOB = "testjob"
+
+BASE_URL_STUDIO_FORMAT = "https://{}.studio.{}.sagemaker.aws/tensorboard/default"
+REDIRECT_STUDIO_FORMAT = (
+    "/data/plugin/sagemaker_data_manager/add_folder_or_job?Redirect=True&Name={}"
+)
+BASE_URL_NON_STUDIO_FORMAT = (
+    "https://{region}.console.aws.amazon.com/sagemaker/home?region={region}#/tensor-board-landing"
+)
+REDIRECT_NON_STUDIO_FORMAT = "/{}"
+
+
+@patch("os.path.isfile")
+def test_tb_init_and_url_non_studio_user(mock_file_exists):
+    """
+    Test TensorBoardApp for non Studio users.
+    """
+    mock_file_exists.return_value = False
+    tb_app = TensorBoardApp(TEST_REGION)
+    assert tb_app.region == TEST_REGION
+    assert tb_app._domain_id is None
+    assert tb_app._user_profile_name is None
+    assert tb_app._valid_domain_and_user is False
+
+    # test url without job redirect
+    assert tb_app.get_app_url() == BASE_URL_NON_STUDIO_FORMAT.format(region=TEST_REGION)
+
+    # test url with valid job redirect
+    assert tb_app.get_app_url(TEST_TRAINING_JOB) == BASE_URL_NON_STUDIO_FORMAT.format(
+        region=TEST_REGION
+    ) + REDIRECT_NON_STUDIO_FORMAT.format(TEST_TRAINING_JOB)
+
+    # test url with invalid job redirect
+    with pytest.raises(ValueError):
+        tb_app.get_app_url("invald_job_name!")
+
+
+@patch("os.path.isfile")
+def test_tb_init_and_url_studio_user_valid_medatada(mock_file_exists):
+    """
+    Test TensorBoardApp for Studio user when the notebook metadata file provided by Studio is valid.
+    """
+    mock_file_exists.return_value = True
+    with patch("builtins.open", mock_open(read_data=TEST_NOTEBOOK_METADATA)):
+        tb_app = TensorBoardApp(TEST_REGION)
+        assert tb_app.region == TEST_REGION
+        assert tb_app._domain_id == TEST_DOMAIN
+        assert tb_app._user_profile_name == TEST_USER_PROFILE
+        assert tb_app._valid_domain_and_user is True
+
+        # test url without job redirect
+        assert (
+            tb_app.get_app_url()
+            == BASE_URL_STUDIO_FORMAT.format(TEST_DOMAIN, TEST_REGION) + "/#sagemaker_data_manager"
+        )
+
+        # test url with valid job redirect
+        assert tb_app.get_app_url(TEST_TRAINING_JOB) == BASE_URL_STUDIO_FORMAT.format(
+            TEST_DOMAIN, TEST_REGION
+        ) + REDIRECT_STUDIO_FORMAT.format(TEST_TRAINING_JOB)
+
+        # test url with invalid job redirect
+        with pytest.raises(ValueError):
+            tb_app.get_app_url("invald_job_name!")
+
+
+@patch("os.path.isfile")
+def test_tb_init_and_url_studio_user_invalid_medatada(mock_file_exists):
+    """
+    Test TensorBoardApp for Studio user when the notebook metadata file provided by Studio is invalid.
+    """
+    mock_file_exists.return_value = True
+
+    # test file does not contain domain and user profle
+    with patch("builtins.open", mock_open(read_data=json.dumps({"Fake": "Fake"}))):
+        assert TensorBoardApp(TEST_REGION).get_app_url() == BASE_URL_NON_STUDIO_FORMAT.format(
+            region=TEST_REGION
+        )
+
+    # test invalid user profile name
+    with patch(
+        "builtins.open",
+        mock_open(read_data=json.dumps({"DomainId": TEST_DOMAIN, "UserProfileName": "u" * 64})),
+    ):
+        assert TensorBoardApp(TEST_REGION).get_app_url() == BASE_URL_NON_STUDIO_FORMAT.format(
+            region=TEST_REGION
+        )
+
+    # test invalid domain id
+    with patch(
+        "builtins.open",
+        mock_open(
+            read_data=json.dumps({"DomainId": "d" * 64, "UserProfileName": TEST_USER_PROFILE})
+        ),
+    ):
+        assert TensorBoardApp(TEST_REGION).get_app_url() == BASE_URL_NON_STUDIO_FORMAT.format(
+            region=TEST_REGION
+        )
+
+
+def test_tb_init_with_default_region():
+    """
+    Test TensorBoardApp init when user does not provide region.
+    """
+    # happy case
+    with patch("sagemaker.Session.boto_region_name", new_callable=PropertyMock) as region_mock:
+        region_mock.return_value = TEST_REGION
+        tb_app = TensorBoardApp()
+        assert tb_app.region == TEST_REGION
+
+    # no default region configured
+    with patch("sagemaker.Session.boto_region_name", new_callable=PropertyMock) as region_mock:
+        region_mock.side_effect = [ValueError()]
+        with pytest.raises(ValueError):
+            tb_app = TensorBoardApp()
