Merge branch 'master' into fix-p2

tejaschumbalkar · web-flow · commit 6dd109c4740a · 2023-04-25T15:34:39.000-07:00
diff --git a/src/sagemaker/image_uri_config/data-wrangler.json b/src/sagemaker/image_uri_config/data-wrangler.json
@@ -7,6 +7,35 @@
           "ap-east-1": "707077482487",
           "ap-northeast-1": "649008135260",
           "ap-northeast-2": "131546521161",
+          "ap-northeast-3": "913387583493",
+          "ap-south-1": "089933028263",
+          "ap-southeast-1": "119527597002",
+          "ap-southeast-2": "422173101802",
+          "ca-central-1": "557239378090",
+          "eu-central-1": "024640144536",
+          "eu-north-1": "054986407534",
+          "eu-south-1": "488287956546",
+          "eu-west-1": "245179582081",
+          "eu-west-2": "894491911112",
+          "eu-west-3": "807237891255",
+          "me-south-1": "376037874950",
+          "sa-east-1": "424196993095",
+          "us-east-1": "663277389841",
+          "us-east-2": "415577184552",
+          "us-west-1": "926135532090",
+          "us-west-2": "174368400705",
+          "cn-north-1": "245909111842",
+          "cn-northwest-1": "249157047649"
+        },
+        "repository": "sagemaker-data-wrangler-container"
+      },
+      "2.x": {
+        "registries": {
+          "af-south-1": "143210264188",
+          "ap-east-1": "707077482487",
+          "ap-northeast-1": "649008135260",
+          "ap-northeast-2": "131546521161",
+          "ap-northeast-3": "913387583493",
           "ap-south-1": "089933028263",
           "ap-southeast-1": "119527597002",
           "ap-southeast-2": "422173101802",
diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
@@ -36,6 +36,7 @@
 SKLEARN_FRAMEWORK = "sklearn"
 TRAINIUM_ALLOWED_FRAMEWORKS = "pytorch"
 INFERENCE_GRAVITON = "inference_graviton"
+DATA_WRANGLER_FRAMEWORK = "data-wrangler"
 
 
 @override_pipeline_parameter_var
@@ -461,6 +462,9 @@ def _validate_version_and_set_if_needed(version, config, framework):
 
         return available_versions[0]
 
+    if version is None and framework in [DATA_WRANGLER_FRAMEWORK]:
+        version = _get_latest_versions(available_versions)
+
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
     return version
 
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/sagemaker/image_uris/test_data_wrangler.py b/tests/unit/sagemaker/image_uris/test_data_wrangler.py
@@ -20,6 +20,7 @@
     "ap-east-1": "707077482487",
     "ap-northeast-1": "649008135260",
     "ap-northeast-2": "131546521161",
+    "ap-northeast-3": "913387583493",
     "ap-south-1": "089933028263",
     "ap-southeast-1": "119527597002",
     "ap-southeast-2": "422173101802",
@@ -39,15 +40,29 @@
     "cn-north-1": "245909111842",
     "cn-northwest-1": "249157047649",
 }
+VERSIONS = ["1.x", "2.x"]
 
 
 def test_data_wrangler_ecr_uri():
-    for region in DATA_WRANGLER_ACCOUNTS.keys():
-        actual_uri = image_uris.retrieve("data-wrangler", region=region)
-        expected_uri = expected_uris.algo_uri(
-            "sagemaker-data-wrangler-container",
-            DATA_WRANGLER_ACCOUNTS[region],
-            region,
-            version="1.x",
-        )
-        assert expected_uri == actual_uri
+    for version in VERSIONS:
+        for region in DATA_WRANGLER_ACCOUNTS.keys():
+            actual_uri = image_uris.retrieve("data-wrangler", region=region, version="1.x")
+            expected_uri = expected_uris.algo_uri(
+                "sagemaker-data-wrangler-container",
+                DATA_WRANGLER_ACCOUNTS[region],
+                region,
+                version="1.x",
+            )
+            assert expected_uri == actual_uri
+
+
+def test_data_wrangler_ecr_uri_none():
+    region = "us-west-2"
+    actual_uri = image_uris.retrieve("data-wrangler", region=region)
+    expected_uri = expected_uris.algo_uri(
+        "sagemaker-data-wrangler-container",
+        DATA_WRANGLER_ACCOUNTS[region],
+        region,
+        version=VERSIONS[-1],
+    )
+    assert expected_uri == actual_uri
diff --git a/tests/unit/sagemaker/wrangler/test_processing.py b/tests/unit/sagemaker/wrangler/test_processing.py
@@ -23,7 +23,7 @@
 REGION = "us-west-2"
 DATA_WRANGLER_RECIPE_SOURCE = "s3://data_wrangler_flows/flow-26-18-43-16-0b48ac2e.flow"
 DATA_WRANGLER_CONTAINER_URI = (
-    "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:1.x"
+    "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:2.x"
 )
 MOCK_S3_URI = "s3://mock_data/mock.csv"
 
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()

Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@`
`23`	`23`	`REGION = "us-west-2"`
`24`	`24`	`DATA_WRANGLER_RECIPE_SOURCE = "s3://data_wrangler_flows/flow-26-18-43-16-0b48ac2e.flow"`
`25`	`25`	`DATA_WRANGLER_CONTAINER_URI = (`
`26`		`- "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:1.x"`
	`26`	`+ "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:2.x"`
`27`	`27`	`)`
`28`	`28`	`MOCK_S3_URI = "s3://mock_data/mock.csv"`
`29`	`29`