AzurePipelinesCredential | adding mlflow uri func (#36580)

* Bug 3323988: Regex fix and indices correction for model download

* fixing test case

* adding mlflow tracking uri func

* passing service context to azureml mlflow

* final flow APC complete

* modify host_url

* fixing unit test cases

* changing mock for urlparse

* fixing the log msg

Author: kshitij-microsoft

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py

@@ -165,7 +165,9 @@ def __init__(
         self.compute_runtime = compute_runtime
 
     @classmethod
-    def _from_rest_object(cls, rest_obj: RestWorkspace) -> Optional["FeatureStore"]:
+    def _from_rest_object(
+        cls, rest_obj: RestWorkspace, v2_service_context: Optional[object] = None
+    ) -> Optional["FeatureStore"]:
         if not rest_obj:
             return None
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/hub.py

@@ -131,11 +131,11 @@ def _get_schema_class(cls):
         return HubSchema
 
     @classmethod
-    def _from_rest_object(cls, rest_obj: RestWorkspace) -> Optional["Hub"]:
+    def _from_rest_object(cls, rest_obj: RestWorkspace, v2_service_context: Optional[object] = None) -> Optional["Hub"]:
         if not rest_obj:
             return None
 
-        workspace_object = Workspace._from_rest_object(rest_obj)
+        workspace_object = Workspace._from_rest_object(rest_obj, v2_service_context)
 
         default_resource_group = None
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py

@@ -314,7 +314,10 @@ def _load(
         return result
 
     @classmethod
-    def _from_rest_object(cls, rest_obj: RestWorkspace) -> Optional["Workspace"]:
+    def _from_rest_object(
+        cls, rest_obj: RestWorkspace, v2_service_context: Optional[object] = None
+    ) -> Optional["Workspace"]:
+
         if not rest_obj:
             return None
         customer_managed_key = (
@@ -329,8 +332,20 @@ def _from_rest_object(cls, rest_obj: RestWorkspace) -> Optional["Workspace"]:
 
         # TODO: Remove attribute check once Oct API version is out
         mlflow_tracking_uri = None
+
         if hasattr(rest_obj, "ml_flow_tracking_uri"):
-            mlflow_tracking_uri = rest_obj.ml_flow_tracking_uri
+            try:
+                from azureml.mlflow import get_mlflow_tracking_uri_v2
+
+                mlflow_tracking_uri = get_mlflow_tracking_uri_v2(rest_obj, v2_service_context)
+            except ImportError:
+                mlflow_tracking_uri = rest_obj.ml_flow_tracking_uri
+                error_msg = (
+                    "azureml.mlflow could not be imported. "
+                    "Please ensure that latest 'azureml-mlflow' has been installed in the current python environment"
+                )
+                print(error_msg)
+                # warnings.warn(error_msg, UserWarning)
 
         # TODO: Remove once Online Endpoints updates API version to at least 2023-08-01
         allow_roleassignment_on_rg = None

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py

@@ -86,11 +86,30 @@ def get(self, workspace_name: Optional[str] = None, **kwargs: Any) -> Optional[W
         workspace_name = self._check_workspace_name(workspace_name)
         resource_group = kwargs.get("resource_group") or self._resource_group_name
         obj = self._operation.get(resource_group, workspace_name)
+        v2_service_context = {}
+
+        v2_service_context["subscription_id"] = self._subscription_id
+        v2_service_context["workspace_name"] = workspace_name
+        v2_service_context["resource_group_name"] = resource_group
+        v2_service_context["auth"] = self._credentials  # type: ignore
+
+        from urllib.parse import urlparse
+
+        if obj is not None and obj.ml_flow_tracking_uri:
+            parsed_url = urlparse(obj.ml_flow_tracking_uri)
+            host_url = "https://{}".format(parsed_url.netloc)
+            v2_service_context["host_url"] = host_url
+        else:
+            v2_service_context["host_url"] = ""
+
+        # host_url=service_context._get_mlflow_url(),
+        # cloud=_get_cloud_or_default(
+        #     service_context.get_auth()._cloud_type.name
         if obj is not None and obj.kind is not None and obj.kind.lower() == WorkspaceKind.HUB:
-            return Hub._from_rest_object(obj)
+            return Hub._from_rest_object(obj, v2_service_context)
         if obj is not None and obj.kind is not None and obj.kind.lower() == WorkspaceKind.PROJECT:
-            return Project._from_rest_object(obj)
-        return Workspace._from_rest_object(obj)
+            return Project._from_rest_object(obj, v2_service_context)
+        return Workspace._from_rest_object(obj, v2_service_context)
 
     def begin_create(
         self,

sdk/ml/azure-ai-ml/tests/workspace/ai_workspaces/unittests/test_mocked_operations.py

@@ -42,7 +42,8 @@ def test_list(self, arg: str, mock_hub_operation: WorkspaceOperations) -> None:
         else:
             mock_hub_operation._operation.list_by_resource_group.assert_called_once()
 
-    def test_get(self, mock_hub_operation: WorkspaceOperations) -> None:
+    def test_get(self, mock_hub_operation: WorkspaceOperations, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         mock_hub_operation.get(name="random_name")
         mock_hub_operation._operation.get.assert_called_once()
 

sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations.py

@@ -1,5 +1,5 @@
 from typing import Optional
-from unittest.mock import ANY, DEFAULT, MagicMock, Mock
+from unittest.mock import ANY, DEFAULT, MagicMock, Mock, patch
 from uuid import UUID, uuid4
 
 import pytest
@@ -20,6 +20,7 @@
 )
 from azure.ai.ml.operations import WorkspaceOperations
 from azure.core.polling import LROPoller
+import urllib.parse
 
 
 @pytest.fixture
@@ -87,7 +88,8 @@ def test_list(self, arg: str, mock_workspace_operation: WorkspaceOperations) ->
         else:
             mock_workspace_operation._operation.list_by_resource_group.assert_called_once()
 
-    def test_get(self, mock_workspace_operation: WorkspaceOperations) -> None:
+    def test_get(self, mock_workspace_operation: WorkspaceOperations, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         mock_workspace_operation.get("random_name")
         mock_workspace_operation._operation.get.assert_called_once()
 
@@ -114,7 +116,8 @@ def test_begin_create(
         mocker.patch("azure.ai.ml._arm_deployments.ArmDeploymentExecutor.deploy_resource", return_value=LROPoller)
         mock_workspace_operation.begin_create(workspace=Workspace(name="name"))
 
-    def test_update(self, mock_workspace_operation: WorkspaceOperations) -> None:
+    def test_update(self, mock_workspace_operation: WorkspaceOperations, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         ws = Workspace(
             name="name",
             description="description",
@@ -135,6 +138,7 @@ def outgoing_call(rg, name, params, polling, cls):
     def test_update_with_role_assignemnt(
         self, mock_workspace_operation: WorkspaceOperations, mocker: MockFixture
     ) -> None:
+        mocker.patch("urllib.parse.urlparse")
         mocker.patch(
             "azure.ai.ml.operations.WorkspaceOperations._populate_feature_store_role_assignment_parameters",
             return_value=({}, {}, {}),
@@ -163,6 +167,7 @@ def outgoing_call(rg, name, params, polling, cls):
         mock_workspace_operation._operation.begin_update.assert_called()
 
     def test_delete(self, mock_workspace_operation: WorkspaceOperations, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         mocker.patch("azure.ai.ml.operations._workspace_operations_base.delete_resource_by_arm_id", return_value=None)
         mocker.patch(
             "azure.ai.ml.operations._workspace_operations_base.get_generic_arm_resource_by_arm_id", return_value=None
@@ -171,6 +176,7 @@ def test_delete(self, mock_workspace_operation: WorkspaceOperations, mocker: Moc
         mock_workspace_operation._operation.begin_delete.assert_called_once()
 
     def test_purge(self, mock_workspace_operation: WorkspaceOperations, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         mocker.patch("azure.ai.ml.operations._workspace_operations_base.delete_resource_by_arm_id", return_value=None)
         mocker.patch(
             "azure.ai.ml.operations._workspace_operations_base.get_generic_arm_resource_by_arm_id", return_value=None

sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py

@@ -31,6 +31,7 @@
 )
 from azure.ai.ml.operations._workspace_operations_base import WorkspaceOperationsBase
 from azure.core.polling import LROPoller
+import urllib.parse
 
 
 @pytest.fixture
@@ -178,6 +179,8 @@ def test_create_get_exception_swallow(
     def test_begin_create_existing_ws(
         self, mock_workspace_operation_base: WorkspaceOperationsBase, mocker: MockFixture
     ):
+        mocker.patch("urllib.parse.urlparse")
+
         def outgoing_call(rg, name, params, polling, cls):
             assert name == "name"
             return DEFAULT
@@ -187,7 +190,8 @@ def outgoing_call(rg, name, params, polling, cls):
         mock_workspace_operation_base.begin_create(workspace=Workspace(name="name"))
         mock_workspace_operation_base._operation.begin_update.assert_called()
 
-    def test_update(self, mock_workspace_operation_base: WorkspaceOperationsBase) -> None:
+    def test_update(self, mock_workspace_operation_base: WorkspaceOperationsBase, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         ws = Workspace(
             name="name",
             tags={"key": "value"},
@@ -244,6 +248,7 @@ def outgoing_call(rg, name, params, polling, cls):
     def test_update_with_empty_property_values(
         self, mock_workspace_operation_base: WorkspaceOperationsBase, mocker: MockFixture
     ) -> None:
+        mocker.patch("urllib.parse.urlparse")
         ws = Workspace(name="name", description="", display_name="", image_build_compute="")
         mocker.patch("azure.ai.ml.operations.WorkspaceOperations.get", return_value=ws)
 
@@ -267,6 +272,7 @@ def outgoing_call(rg, name, params, polling, cls):
         mock_workspace_operation_base._operation.begin_update.assert_called()
 
     def test_delete_no_wait(self, mock_workspace_operation_base: WorkspaceOperationsBase, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         mocker.patch("azure.ai.ml.operations._workspace_operations_base.delete_resource_by_arm_id", return_value=None)
         mocker.patch(
             "azure.ai.ml.operations._workspace_operations_base.get_generic_arm_resource_by_arm_id", return_value=None
@@ -275,6 +281,7 @@ def test_delete_no_wait(self, mock_workspace_operation_base: WorkspaceOperations
         mock_workspace_operation_base._operation.begin_delete.assert_called_once()
 
     def test_delete_wait(self, mock_workspace_operation_base: WorkspaceOperationsBase, mocker: MockFixture) -> None:
+        mocker.patch("urllib.parse.urlparse")
         mocker.patch("azure.ai.ml.operations._workspace_operations_base.delete_resource_by_arm_id", return_value=None)
         mocker.patch(
             "azure.ai.ml.operations._workspace_operations_base.get_generic_arm_resource_by_arm_id", return_value=None
@@ -600,6 +607,7 @@ def test_update_workspace_with_serverless_custom_vnet(
         mock_workspace_operation_base: WorkspaceOperationsBase,
         mocker: MockFixture,
     ) -> None:
+        mocker.patch("urllib.parse.urlparse")
         ws = Workspace(name="name", location="test", serverless_compute=serverless_compute_settings)
         spy = mocker.spy(mock_workspace_operation_base._operation, "begin_update")
         mock_workspace_operation_base.begin_update(ws)