Skip to content

AzurePipelinesCredential | adding mlflow uri func #36580

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 17 commits into from
Oct 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
648d0a9
Bug 3323988: Regex fix and indices correction for model download
kshitij-microsoft Jul 2, 2024
c21f605
Merge branch 'main' of github.com:Azure/azure-sdk-for-python
kshitij-microsoft Jul 4, 2024
331d7e3
Merge branch 'main' of github.com:Azure/azure-sdk-for-python
kshitij-microsoft Jul 5, 2024
f125053
fixing test case
kshitij-microsoft Jul 17, 2024
a5aa64a
Merge branch 'main' of github.com:Azure/azure-sdk-for-python
kshitij-microsoft Jul 23, 2024
921e805
adding mlflow tracking uri func
kshitij-microsoft Jul 23, 2024
d4785cd
Merge branch 'main' of github.com:Azure/azure-sdk-for-python
kshitij-microsoft Jul 31, 2024
e647204
passing service context to azureml mlflow
kshitij-microsoft Aug 2, 2024
2ed6544
Merge branch 'main' of github.com:Azure/azure-sdk-for-python
kshitij-microsoft Aug 28, 2024
a03c7a3
final flow APC complete
kshitij-microsoft Sep 5, 2024
357698f
Merge branch 'main' of github.com:Azure/azure-sdk-for-python
kshitij-microsoft Sep 6, 2024
d688c3c
Merge branch 'main' of https://github.com/Azure/azure-sdk-for-python
kshitij-microsoft Sep 24, 2024
3dbb2d2
modify host_url
kshitij-microsoft Sep 25, 2024
58a96dc
fixing unit test cases
kshitij-microsoft Sep 27, 2024
ae32d9d
Merge branch 'main' of https://github.com/Azure/azure-sdk-for-python
kshitij-microsoft Sep 30, 2024
0ab7bc4
changing mock for urlparse
kshitij-microsoft Oct 1, 2024
5a90da5
fixing the log msg
kshitij-microsoft Oct 24, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,9 @@ def __init__(
self.managed_network = managed_network

@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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
19 changes: 17 additions & 2 deletions sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -20,6 +20,7 @@
)
from azure.ai.ml.operations import WorkspaceOperations
from azure.core.polling import LROPoller
import urllib.parse


@pytest.fixture
Expand Down Expand Up @@ -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()

Expand All @@ -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",
Expand All @@ -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=({}, {}, {}),
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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"},
Expand Down Expand Up @@ -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)

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Loading