feat: (sagemaker-mlflow) Adding Presigned Url function to SDK (aws#1462) (aws#1477)

* mlflow presigned url changes

* addressing design feedback

* test changes

Author: ananth102

src/sagemaker/mlflow/__init__.py

@@ -0,0 +1,12 @@
+# 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.

src/sagemaker/mlflow/tracking_server.py

@@ -0,0 +1,50 @@
+# 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 code related to the Mlflow Tracking Server."""
+
+from __future__ import absolute_import
+from typing import Optional, TYPE_CHECKING
+from sagemaker.apiutils import _utils
+
+if TYPE_CHECKING:
+    from sagemaker import Session
+
+
+def generate_mlflow_presigned_url(
+    name: str,
+    expires_in_seconds: Optional[int] = None,
+    session_expiration_duration_in_seconds: Optional[int] = None,
+    sagemaker_session: Optional["Session"] = None,
+) -> str:
+    """Generate a presigned url to acess the Mlflow UI.
+
+    Args:
+        name (str): Name of the Mlflow Tracking Server
+        expires_in_seconds (int): Expiration time of the first usage
+               of the presigned url in seconds.
+        session_expiration_duration_in_seconds (int): Session duration of the presigned url in
+               seconds after the first use.
+        sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed. If not specified, one is created using the
+                default AWS configuration chain.
+    Returns:
+        (str): Authorized Url to acess the Mlflow UI.
+    """
+    session = sagemaker_session or _utils.default_session()
+    api_response = session.create_presigned_mlflow_tracking_server_url(
+        name, expires_in_seconds, session_expiration_duration_in_seconds
+    )
+    return api_response["AuthorizedUrl"]

src/sagemaker/session.py

@@ -6705,6 +6705,36 @@ def wait_for_inference_recommendations_job(
         _check_job_status(job_name, desc, "Status")
         return desc
 
+    def create_presigned_mlflow_tracking_server_url(
+        self,
+        tracking_server_name: str,
+        expires_in_seconds: int = None,
+        session_expiration_duration_in_seconds: int = None,
+    ) -> Dict[str, Any]:
+        """Creates a Presigned Url to acess the Mlflow UI.
+
+        Args:
+            tracking_server_name (str): Name of the Mlflow Tracking Server.
+            expires_in_seconds (int): Expiration duration of the URL.
+            session_expiration_duration_in_seconds (int): Session duration of the URL.
+        Returns:
+            (dict): Return value from the ``CreatePresignedMlflowTrackingServerUrl`` API.
+
+        """
+
+        create_presigned_url_args = {"TrackingServerName": tracking_server_name}
+        if expires_in_seconds is not None:
+            create_presigned_url_args["ExpiresInSeconds"] = expires_in_seconds
+
+        if session_expiration_duration_in_seconds is not None:
+            create_presigned_url_args["SessionExpirationDurationInSeconds"] = (
+                session_expiration_duration_in_seconds
+            )
+
+        return self.sagemaker_client.create_presigned_mlflow_tracking_server_url(
+            **create_presigned_url_args
+        )
+
 
 def get_model_package_args(
     content_types=None,

tests/unit/sagemaker/mlflow/__init__.py

@@ -0,0 +1,42 @@
+# 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.mlflow.tracking_server import generate_mlflow_presigned_url
+
+
+def test_generate_presigned_url(sagemaker_session):
+    client = sagemaker_session.sagemaker_client
+    client.create_presigned_mlflow_tracking_server_url.return_value = {
+        "AuthorizedUrl": "https://t-wo.example.com",
+    }
+    url = generate_mlflow_presigned_url(
+        "w",
+        expires_in_seconds=10,
+        session_expiration_duration_in_seconds=5,
+        sagemaker_session=sagemaker_session,
+    )
+    client.create_presigned_mlflow_tracking_server_url.assert_called_with(
+        TrackingServerName="w", ExpiresInSeconds=10, SessionExpirationDurationInSeconds=5
+    )
+    assert url == "https://t-wo.example.com"
+
+
+def test_generate_presigned_url_minimal(sagemaker_session):
+    client = sagemaker_session.sagemaker_client
+    client.create_presigned_mlflow_tracking_server_url.return_value = {
+        "AuthorizedUrl": "https://t-wo.example.com",
+    }
+    url = generate_mlflow_presigned_url("w", sagemaker_session=sagemaker_session)
+    client.create_presigned_mlflow_tracking_server_url.assert_called_with(TrackingServerName="w")
+    assert url == "https://t-wo.example.com"

tests/unit/sagemaker/mlflow/test_tracking_server.py

@@ -6263,6 +6263,24 @@ def test_create_inference_recommendations_job_propogate_other_exception(
     assert "AccessDeniedException" in str(error)
 
 
+def test_create_presigned_mlflow_tracking_server_url(sagemaker_session):
+    sagemaker_session.create_presigned_mlflow_tracking_server_url("ts", 1, 2)
+    assert (
+        sagemaker_session.sagemaker_client.create_presigned_mlflow_tracking_server_url.called_with(
+            TrackingServerName="ts", ExpiresInSeconds=1, SessionExpirationDurationInSeconds=2
+        )
+    )
+
+
+def test_create_presigned_mlflow_tracking_server_url_minimal(sagemaker_session):
+    sagemaker_session.create_presigned_mlflow_tracking_server_url("ts")
+    assert (
+        sagemaker_session.sagemaker_client.create_presigned_mlflow_tracking_server_url.called_with(
+            TrackingServerName="ts"
+        )
+    )
+
+
 DEFAULT_LOG_EVENTS_INFERENCE_RECOMMENDER = [
     MockBotoException("ResourceNotFoundException"),
     {"nextForwardToken": None, "events": [{"timestamp": 1, "message": "hi there #1"}]},