1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
|
# ---------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# ---------------------------------------------------------
import logging
from typing import Optional, Tuple
from typing_extensions import Literal
from azure.ai.ml._azure_environments import _get_default_cloud_name, _get_registry_discovery_endpoint_from_metadata
from azure.ai.ml._restclient.registry_discovery import AzureMachineLearningWorkspaces as ServiceClientRegistryDiscovery
from azure.ai.ml._restclient.v2021_10_01_dataplanepreview import AzureMachineLearningWorkspaces
from azure.ai.ml._restclient.v2021_10_01_dataplanepreview.models import (
BlobReferenceSASRequestDto,
TemporaryDataReferenceRequestDto,
)
from azure.ai.ml.constants._common import REGISTRY_ASSET_ID
from azure.ai.ml.exceptions import MlException
from azure.core.exceptions import HttpResponseError
module_logger = logging.getLogger(__name__)
MFE_PATH_PREFIX = "mferp/managementfrontend"
class RegistryDiscovery:
def __init__(
self,
credential: str,
registry_name: str,
service_client_registry_discovery_client: ServiceClientRegistryDiscovery,
**kwargs,
):
self.credential = credential
self.registry_name = registry_name
self.service_client_registry_discovery_client = service_client_registry_discovery_client
self.kwargs = kwargs
self._resource_group = None
self._subscription_id = None
self._base_url = None
self.workspace_region = kwargs.get("workspace_location", None)
def _get_registry_details(self) -> str:
response = self.service_client_registry_discovery_client.registry_management_non_workspace.registry_management_non_workspace( # pylint: disable=line-too-long
self.registry_name
)
if self.workspace_region:
_check_region_fqdn(self.workspace_region, response)
self._base_url = f"https://cert-{self.workspace_region}.experiments.azureml.net/{MFE_PATH_PREFIX}"
else:
self._base_url = f"{response.primary_region_resource_provider_uri}{MFE_PATH_PREFIX}"
self._subscription_id = response.subscription_id
self._resource_group = response.resource_group
def get_registry_service_client(self) -> AzureMachineLearningWorkspaces:
self._get_registry_details()
self.kwargs.pop("subscription_id", None)
self.kwargs.pop("resource_group", None)
service_client_10_2021_dataplanepreview = AzureMachineLearningWorkspaces(
subscription_id=self._subscription_id,
resource_group=self._resource_group,
credential=self.credential,
base_url=self._base_url,
**self.kwargs,
)
return service_client_10_2021_dataplanepreview
@property
def subscription_id(self) -> str:
"""The subscription id of the registry.
:return: Subscription Id
:rtype: str
"""
return self._subscription_id
@property
def resource_group(self) -> str:
"""The resource group of the registry.
:return: Resource Group
:rtype: str
"""
return self._resource_group
def get_sas_uri_for_registry_asset(service_client, name, version, resource_group, registry, body) -> str:
"""Get sas_uri for registry asset.
:param service_client: Service client
:type service_client: AzureMachineLearningWorkspaces
:param name: Asset name
:type name: str
:param version: Asset version
:type version: str
:param resource_group: Resource group
:type resource_group: str
:param registry: Registry name
:type registry: str
:param body: Request body
:type body: TemporaryDataReferenceRequestDto
:rtype: str
"""
sas_uri = None
try:
res = service_client.temporary_data_references.create_or_get_temporary_data_reference(
name=name,
version=version,
resource_group_name=resource_group,
registry_name=registry,
body=body,
)
sas_uri = res.blob_reference_for_consumption.credential.additional_properties["sasUri"]
except HttpResponseError as e:
# "Asset already exists" exception is thrown from service with error code 409, that we need to ignore
if e.status_code == 409:
module_logger.debug("Skipping file upload, reason: %s", str(e.reason))
else:
raise e
return sas_uri
def get_asset_body_for_registry_storage(
registry_name: str, asset_type: str, asset_name: str, asset_version: str
) -> TemporaryDataReferenceRequestDto:
"""Get Asset body for registry.
:param registry_name: Registry name.
:type registry_name: str
:param asset_type: Asset type.
:type asset_type: str
:param asset_name: Asset name.
:type asset_name: str
:param asset_version: Asset version.
:type asset_version: str
:return: The temporary data reference request dto
:rtype: TemporaryDataReferenceRequestDto
"""
body = TemporaryDataReferenceRequestDto(
asset_id=REGISTRY_ASSET_ID.format(registry_name, asset_type, asset_name, asset_version),
temporary_data_reference_type="TemporaryBlobReference",
)
return body
def get_storage_details_for_registry_assets(
service_client: AzureMachineLearningWorkspaces,
asset_type: str,
asset_name: str,
asset_version: str,
rg_name: str,
reg_name: str,
uri: str,
) -> Tuple[str, Literal["NoCredentials", "SAS"]]:
"""Get storage details for registry assets.
:param service_client: AzureMachineLearningWorkspaces service client.
:type service_client: AzureMachineLearningWorkspaces
:param asset_type: Asset type.
:type asset_type: str
:param asset_name: Asset name.
:type asset_name: str
:param asset_version: Asset version.
:type asset_version: str
:param rg_name: Resource group name.
:type rg_name: str
:param reg_name: Registry name
:type reg_name: str
:param uri: asset uri
:type uri: str
:return: A 2-tuple of a URI and a string. Either:
* A blob uri and "NoCredentials"
* A sas URI and "SAS"
:rtype: Tuple[str, Literal["NoCredentials", "SAS"]]
"""
body = BlobReferenceSASRequestDto(
asset_id=REGISTRY_ASSET_ID.format(reg_name, asset_type, asset_name, asset_version),
blob_uri=uri,
)
sas_uri = service_client.data_references.get_blob_reference_sas(
name=asset_name,
version=asset_version,
resource_group_name=rg_name,
registry_name=reg_name,
body=body,
)
if sas_uri.blob_reference_for_consumption.credential.credential_type == "no_credentials":
return sas_uri.blob_reference_for_consumption.blob_uri, "NoCredentials"
return (
sas_uri.blob_reference_for_consumption.credential.additional_properties["sasUri"],
"SAS",
)
def get_registry_client(credential, registry_name, workspace_location: Optional[str] = None, **kwargs):
base_url = _get_registry_discovery_endpoint_from_metadata(_get_default_cloud_name())
kwargs.pop("base_url", None)
service_client_registry_discovery_client = ServiceClientRegistryDiscovery(
credential=credential, base_url=base_url, **kwargs
)
if workspace_location:
workspace_kwargs = {"workspace_location": workspace_location}
kwargs.update(workspace_kwargs)
registry_discovery = RegistryDiscovery(
credential, registry_name, service_client_registry_discovery_client, **kwargs
)
service_client_10_2021_dataplanepreview = registry_discovery.get_registry_service_client()
subscription_id = registry_discovery.subscription_id
resource_group_name = registry_discovery.resource_group
return service_client_10_2021_dataplanepreview, resource_group_name, subscription_id
def _check_region_fqdn(workspace_region, response):
if workspace_region in response.additional_properties["registryFqdns"].keys():
return
regions = list(response.additional_properties["registryFqdns"].keys())
msg = f"Workspace region {workspace_region} not supported by the \
registry {response.registry_name} regions {regions}"
raise MlException(message=msg, no_personal_data_message=msg)
|
