feat: Enable notebook instances to get presigned url (#4107)

Author: mseth10

src/sagemaker/interactive_apps/base_interactive_app.py

@@ -59,6 +59,7 @@ def __init__(
         # Used to store domain and user profile info retrieved from Studio environment.
         self._domain_id = None
         self._user_profile_name = None
+        self._in_studio_env = False
         self._get_domain_and_user()
 
     def __str__(self):
@@ -70,25 +71,25 @@ def __repr__(self):
         return self.__str__()
 
     def _get_domain_and_user(self):
-        """Get and validate studio domain id and user profile from studio environment."""
-        if not self._is_in_studio():
+        """Get domain id and user profile from Studio environment.
+
+        To verify Studio environment, we check if NOTEBOOK_METADATA_FILE exists
+        and domain id and user profile name are present in the file.
+        """
+        if not os.path.isfile(NOTEBOOK_METADATA_FILE):
             return
 
         try:
             with open(NOTEBOOK_METADATA_FILE, "rb") as metadata_file:
                 metadata = json.loads(metadata_file.read())
-                if not self._validate_domain_id(
-                    metadata.get("DomainId")
-                ) or not self._validate_user_profile_name(metadata.get("UserProfileName")):
-                    logger.warning(
-                        "NOTEBOOK_METADATA_FILE detected but failed to get valid domain and user"
-                        " from it."
-                    )
-                    return
-                self._domain_id = metadata.get("DomainId")
-                self._user_profile_name = metadata.get("UserProfileName")
         except OSError as err:
-            logger.warning("Could not load Studio metadata due to unexpected error. %s", err)
+            logger.warning("Could not load metadata due to unexpected error. %s", err)
+            return
+
+        if "DomainId" in metadata and "UserProfileName" in metadata:
+            self._in_studio_env = True
+            self._domain_id = metadata.get("DomainId")
+            self._user_profile_name = metadata.get("UserProfileName")
 
     def _get_presigned_url(
         self,
@@ -142,10 +143,6 @@ def _get_presigned_url(
 
         return url
 
-    def _is_in_studio(self):
-        """Check to see if NOTEBOOK_METADATA_FILE exists to verify Studio environment."""
-        return os.path.isfile(NOTEBOOK_METADATA_FILE)
-
     def _open_url_in_web_browser(self, url: str):
         """Open a URL in the default web browser.
 
@@ -154,23 +151,6 @@ def _open_url_in_web_browser(self, url: str):
         """
         webbrowser.open(url)
 
-    def _validate_domain_id(self, domain_id: Optional[str] = None):
-        """Validate domain id format.
-
-        Args:
-            domain_id (str): Optional. The domain ID to validate. If one is not supplied,
-            self._domain_id will be used instead.
-            Default: ``None``
-
-        Returns:
-            bool: Whether the supplied domain ID is valid.
-        """
-        if domain_id is None:
-            domain_id = self._domain_id
-        if domain_id is None or len(domain_id) > 63:
-            return False
-        return True
-
     def _validate_job_name(self, job_name: str):
         """Validate training job name format.
 
@@ -186,30 +166,35 @@ def _validate_job_name(self, job_name: str):
                 f"Invalid job name. Job name must match regular expression {job_name_regex}"
             )
 
-    def _validate_user_profile_name(self, user_profile_name: Optional[str] = None):
+    def _validate_domain_id(self, domain_id: str):
+        """Validate domain id format.
+
+        Args:
+            domain_id (str): Required. The domain ID to validate.
+
+        Returns:
+            bool: Whether the supplied domain ID is valid.
+        """
+        if domain_id is None or len(domain_id) > 63:
+            return False
+        return True
+
+    def _validate_user_profile_name(self, user_profile_name: str):
         """Validate user profile name format.
 
         Args:
-            user_profile_name (str): Optional. The user profile name to validate. If one is not
-            supplied, self._user_profile_name will be used instead.
-            Default: ``None``
+            user_profile_name (str): Required. The user profile name to validate.
 
         Returns:
             bool: Whether the supplied user profile name is valid.
         """
-        if user_profile_name is None:
-            user_profile_name = self._user_profile_name
         user_profile_name_regex = "^[a-zA-Z0-9](-*[a-zA-Z0-9]){0,62}"
         if user_profile_name is None or not re.fullmatch(
             user_profile_name_regex, user_profile_name
         ):
             return False
         return True
 
-    def _validate_domain_and_user(self):
-        """Helper function to consolidate validation calls."""
-        return self._validate_domain_id() and self._validate_user_profile_name()
-
     @abc.abstractmethod
     def get_app_url(self):
         """Abstract method to generate a URL to help access the application in Studio.

src/sagemaker/interactive_apps/tensorboard.py

@@ -87,36 +87,11 @@ def get_app_url(
         if training_job_name is not None:
             self._validate_job_name(training_job_name)
 
-        if optional_create_presigned_url_kwargs is None:
-            optional_create_presigned_url_kwargs = {}
-
-        if domain_id is not None:
-            optional_create_presigned_url_kwargs["DomainId"] = domain_id
-
-        if user_profile_name is not None:
-            optional_create_presigned_url_kwargs["UserProfileName"] = user_profile_name
-
         if (
-            create_presigned_domain_url
-            and not self._is_in_studio()
-            and self._validate_domain_id(optional_create_presigned_url_kwargs.get("DomainId"))
-            and self._validate_user_profile_name(
-                optional_create_presigned_url_kwargs.get("UserProfileName")
-            )
+            self._in_studio_env
+            and self._validate_domain_id(self._domain_id)
+            and self._validate_user_profile_name(self._user_profile_name)
         ):
-            state_to_encode = None
-            redirect = "TensorBoard"
-
-            if training_job_name is not None:
-                state_to_encode = (
-                    "/tensorboard/default/data/plugin/sagemaker_data_manager/"
-                    + f"add_folder_or_job?Redirect=True&Name={training_job_name}"
-                )
-
-            url = self._get_presigned_url(
-                optional_create_presigned_url_kwargs, redirect, state_to_encode
-            )
-        elif self._is_in_studio() and self._validate_domain_and_user():
             if domain_id or user_profile_name:
                 logger.warning(
                     "Ignoring passed in domain_id and user_profile_name for Studio set values."
@@ -126,15 +101,37 @@ def get_app_url(
                 + "sagemaker.aws/tensorboard/default"
             )
             if training_job_name is not None:
-                self._validate_job_name(training_job_name)
                 url += (
                     "/data/plugin/sagemaker_data_manager/"
                     + f"add_folder_or_job?Redirect=True&Name={training_job_name}"
                 )
             else:
                 url += "/#sagemaker_data_manager"
+
+        elif (
+            not self._in_studio_env
+            and create_presigned_domain_url
+            and self._validate_domain_id(domain_id)
+            and self._validate_user_profile_name(user_profile_name)
+        ):
+            if optional_create_presigned_url_kwargs is None:
+                optional_create_presigned_url_kwargs = {}
+            optional_create_presigned_url_kwargs["DomainId"] = domain_id
+            optional_create_presigned_url_kwargs["UserProfileName"] = user_profile_name
+
+            redirect = "TensorBoard"
+            state_to_encode = None
+            if training_job_name is not None:
+                state_to_encode = (
+                    "/tensorboard/default/data/plugin/sagemaker_data_manager/"
+                    + f"add_folder_or_job?Redirect=True&Name={training_job_name}"
+                )
+
+            url = self._get_presigned_url(
+                optional_create_presigned_url_kwargs, redirect, state_to_encode
+            )
         else:
-            if domain_id or user_profile_name or create_presigned_domain_url:
+            if not self._in_studio_env and create_presigned_domain_url:
                 logger.warning(
                     "A valid domain ID and user profile name were not provided. "
                     "Providing default landing page URL as a result."